Update

[ghstack-poisoned]
pytorch · Oct 3, 2024 · 9798b2d · 9798b2d
1 parent 009abc3
commit 9798b2d
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tensordict/base.py b/tensordict/base.py
@@ -10417,11 +10417,13 @@ def set_(x):
         result._consolidated = {"storage": storage_cast}
         if "metadata" in self._consolidated:
             result._consolidated["metadata"] = deepcopy(self._consolidated["metadata"])
-        if not non_blocking:
-            if device.type == "cuda":
+        if non_blocking in (False, None):
+            if device.type == "cuda" and non_blocking is False:
+                # sending to CUDA force sync
                 cuda_device = device
             elif storage.device.type == "cuda":
-                cuda_device = device
+                # sending from cuda: need sync unless intentionally not asked for
+                cuda_device = storage.device.type
             else:
                 cuda_device = None
             if cuda_device is not None: