amend

torchrl/data/replay_buffers/storages.py

@@ -1280,11 +1280,6 @@ def __repr__(self):
 
 # Utils
 def _mem_map_tensor_as_tensor(mem_map_tensor) -> torch.Tensor:
-    if _CKPT_BACKEND == "torchsnapshot" and not _has_ts:
-        raise ImportError(
-            "the checkpointing backend is set to torchsnapshot but the library is not installed. Consider installing the library or switch to another backend. "
-            f"Supported backends are {_CKPT_BACKEND.backends}"
-        )
     if isinstance(mem_map_tensor, torch.Tensor):
         # This will account for MemoryMappedTensors
         return mem_map_tensor

torchrl/envs/model_based/common.py

@@ -5,9 +5,8 @@
 
 import abc
 import warnings
-from typing import List, Optional, Union
+from typing import List, Optional
 
-import numpy as np
 import torch
 from tensordict import TensorDict
 from tensordict.nn import TensorDictModule

torchrl/envs/model_based/dreamer.py

@@ -3,9 +3,8 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple
 
-import numpy as np
 import torch
 from tensordict import TensorDict
 from tensordict.nn import TensorDictModule