From 1d13105a75392c81f360549ad57c126205716fb1 Mon Sep 17 00:00:00 2001 From: vmoens Date: Fri, 20 Oct 2023 17:07:52 -0700 Subject: [PATCH] amend --- test/test_tensordict.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/test_tensordict.py b/test/test_tensordict.py index 832a87001..bc5d0d19b 100644 --- a/test/test_tensordict.py +++ b/test/test_tensordict.py @@ -2994,6 +2994,8 @@ def stacked_td(self, device, dtype): return stack_td([td1, td2], 2) def memmap_td(self, device, dtype): + if device.type != "cpu": + pytest.skip("MemoryMappedTensors can only be placed on CPU.") return self.td(device, dtype).memmap_() def share_memory_td(self, device, dtype): @@ -3027,7 +3029,7 @@ def test_repr_memmap(self, device, dtype): fields={{ a: MemoryMappedTensor(shape=torch.Size([4, 3, 2, 1, 5]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor})}}, batch_size=torch.Size([4, 3, 2, 1]), - device={str(device)}, + device=cpu, is_shared={is_shared})""" assert repr(tensordict) == expected