non_blocking=True

pytorch · Jan 5, 2024 · 661886d · 661886d
1 parent 24914c8
commit 661886d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensordict/memmap.py b/tensordict/memmap.py
@@ -204,7 +204,7 @@ def from_tensor(
         if copy_data:
             if isinstance(input, DTensor):
                 input = input.full_tensor()
-            out.copy_(input)
+            out.copy_(input, non_blocking=True)
         return out
 
     @property