add torch interface numpy path device identify

Author: refraction-ray

CHANGELOG.md

@@ -28,6 +28,8 @@
 
 - Fixed `TorchLayer` parameter list auto registeration
 
+- Pytorch interface is now device aware (#25)
+
 ## 0.2.1
 
 ### Added

tensorcircuit/interfaces/torch.py

@@ -70,7 +70,8 @@ def forward(ctx: Any, *x: Any) -> Any:  # type: ignore
             # (x, )
             if len(ctx.xdtype) == 1:
                 ctx.xdtype = ctx.xdtype[0]
-
+            if not enable_dlpack:
+                ctx.device = (backend.tree_flatten(x)[0][0]).device
             x = general_args_to_backend(x, enable_dlpack=enable_dlpack)
             y = fun(*x)
             ctx.ydtype = backend.tree_map(lambda s: s.dtype, y)
@@ -80,6 +81,8 @@ def forward(ctx: Any, *x: Any) -> Any:  # type: ignore
             y = general_args_to_backend(
                 y, target_backend="pytorch", enable_dlpack=enable_dlpack
             )
+            if not enable_dlpack:
+                y = backend.tree_map(lambda s: s.to(device=ctx.device), y)
             return y
 
         @staticmethod
@@ -101,6 +104,8 @@ def backward(ctx: Any, *grad_y: Any) -> Any:
                 target_backend="pytorch",
                 enable_dlpack=enable_dlpack,
             )
+            if not enable_dlpack:
+                r = backend.tree_map(lambda s: s.to(device=ctx.device), r)
             if not is_sequence(r):
                 return (r,)
             return r