allow device move even with dlpack on torch interface

Author: refraction-ray

examples/hybrid_gpu_pipeline.py

@@ -25,7 +25,11 @@
 enable_dlpack = True
 # enable_dlpack = False # for old version of ML libs
 tf_device = "/GPU:0"
-# tf_device "/device:CPU:0"
+# tf_device = "/device:CPU:0"
+# another scheme to globally close GPU only for tf
+# https://datascience.stackexchange.com/a/76039
+# but if gpu support is fully shut down as above
+# dlpack=True wont work
 
 # dataset preparation
 

tensorcircuit/interfaces/torch.py

@@ -70,8 +70,7 @@ 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
+            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)
@@ -81,8 +80,7 @@ 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)
+            y = backend.tree_map(lambda s: s.to(device=ctx.device), y)
             return y
 
         @staticmethod
@@ -104,8 +102,7 @@ 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)
+            r = backend.tree_map(lambda s: s.to(device=ctx.device), r)
             if not is_sequence(r):
                 return (r,)
             return r