Update

[ghstack-poisoned]

test/test_tensordict.py

@@ -132,6 +132,17 @@
 mp_ctx = "fork" if (not torch.cuda.is_available() and not _IS_WINDOWS) else "spawn"
 
 
+@pytest.fixture
+def device_fixture():
+    device = torch.get_default_device()
+    if torch.cuda.is_available():
+        torch.set_default_device(torch.device("cuda:0"))
+    elif torch.backends.mps.is_available():
+        torch.set_default_device(torch.device("mps:0"))
+    yield
+    torch.set_default_device(device)
+
+
 def _compare_tensors_identity(td0, td1):
     if isinstance(td0, LazyStackedTensorDict):
         if not isinstance(td1, LazyStackedTensorDict):
@@ -243,7 +254,7 @@ def test_batchsize_reset(self):
             td_u.to_tensordict().batch_size = [1]
 
     @pytest.mark.parametrize("count_duplicates", [False, True])
-    def test_bytes(self, count_duplicates):
+    def test_bytes(self, count_duplicates, device_fixture):
         tensor = torch.zeros(3)
         tensor_with_grad = torch.ones(3, requires_grad=True)
         (tensor_with_grad + 1).sum().backward()
@@ -267,7 +278,7 @@ def test_bytes(self, count_duplicates):
         else:
             assert td.bytes(count_duplicates=count_duplicates) == 12 + 24 + 52 + 0
 
-    def test_depth(ggself):
+    def test_depth(self):
         td = TensorDict({"a": {"b": {"c": {"d": 0}, "e": 0}, "f": 0}, "g": 0}).lock_()
         assert td.depth == 3
         with td.unlock_():