amend

pytorch · Oct 19, 2023 · 3e31297 · 3e31297
1 parent af7f289
commit 3e31297
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensordict/memmap_refact.py b/tensordict/memmap_refact.py
@@ -163,10 +163,12 @@ def from_handler(cls, handler, dtype, shape, index):
         return out
 
     def __setstate__(self, state):
+        print('setstate', self.__dict__)
         if 'filename' in state:
             self.__dict__.update(type(self).from_filename(**state).__dict__)
         else:
             self.__dict__.update(type(self).from_handler(**state).__dict__)
+        assert self.is_shared()
 
     def __getstate__(self):
         if getattr(self, "_handler", None) is not None: