Update

[ghstack-poisoned]

Author: Vincent Moens

tensordict/_td.py

@@ -2770,6 +2770,8 @@ def _memmap_(
         if inplace:
             self._is_memmap = True
             self._is_shared = False  # since they are mutually exclusive
+            if hasattr(self, "_validate_value_cached"):
+                delattr(self, "_validate_value_cached")
             self._device = torch.device("cpu")
         else:
             dest._is_memmap = True

tensordict/base.py

@@ -4544,13 +4544,17 @@ def clear_device_(self) -> T:
 
         """
         self._device = None
+        if hasattr(self, "_validate_value_cached"):
+            delattr(self, "_validate_value_cached")
         for value in self.values():
             if _is_tensor_collection(type(value)):
                 value.clear_device_()
         return self
 
     def _set_device(self, device: torch.device) -> T:
         self._device = device
+        if hasattr(self, "_validate_value_cached"):
+            delattr(self, "_validate_value_cached")
         for value in self.values():
             if _is_tensor_collection(type(value)):
                 value._set_device(device=device)