[BugFix] Fix MPS sync in device transform (#2061)

pytorch · Apr 7, 2024 · 4488c25 · 4488c25
1 parent f85da4c
commit 4488c25
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -3805,7 +3805,7 @@ def __init__(
             if torch.cuda.is_available():
                 self._sync_device = torch.cuda.synchronize
             elif torch.backends.mps.is_available():
-                self._sync_device = torch.cuda.synchronize
+                self._sync_device = torch.mps.synchronize
             elif device.type == "cpu":
                 self._sync_device = _do_nothing
         else:
@@ -3901,7 +3901,7 @@ def _sync_orig_device(self):
                 if torch.cuda.is_available():
                     self._sync_orig_device_val = torch.cuda.synchronize
                 elif torch.backends.mps.is_available():
-                    self._sync_orig_device_val = torch.cuda.synchronize
+                    self._sync_orig_device_val = torch.mps.synchronize
                 elif device.type == "cpu":
                     self._sync_orig_device_val = _do_nothing
             else: