amend

pytorch · Oct 19, 2023 · de0c028 · de0c028
1 parent 7d174a4
commit de0c028
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
@@ -4741,8 +4741,6 @@ def memmap_(
     def load_memmap(cls, prefix: str) -> T:
         prefix = Path(prefix)
         metadata = torch.load(prefix / "meta.pt")
-        # TODO: remove this
-        assert metadata["device"] == torch.device("cpu"), metadata
         out = cls({}, batch_size=metadata["batch_size"], device=metadata["device"])
 
         for path in prefix.glob("**/*meta.pt"):