[Feature] Replaced device_safe() with device

test/test_tensordict.py

@@ -70,13 +70,11 @@ def test_tensordict_set(device):
 @pytest.mark.parametrize("device", get_available_devices())
 def test_tensordict_device(device):
  tensordict = TensorDict({"a": torch.randn(3, 4)}, [])
- with pytest.raises(RuntimeError):
- tensordict.device
+ assert tensordict.device is None
 
  tensordict = TensorDict({"a": torch.randn(3, 4, device=device)}, [])
  assert tensordict["a"].device == device
- with pytest.raises(RuntimeError):
- tensordict.device
+ assert tensordict.device is None
 
  tensordict = TensorDict(
  {
@@ -2104,12 +2102,10 @@ def test_create_on_device():
 
  # TensorDict
  td = TensorDict({}, [5])
- with pytest.raises(RuntimeError):
- td.device
+ assert td.device is None
 
  td.set("a", torch.randn(5, device=device))
- with pytest.raises(RuntimeError):
- td.device
+ assert td.device is None
 
  td = TensorDict({}, [5], device="cuda:0")
  td.set("a", torch.randn(5, 1))
@@ -2119,11 +2115,10 @@ def test_create_on_device():
  td1 = TensorDict({}, [5])
  td2 = TensorDict({}, [5])
  stackedtd = stack_td([td1, td2], 0)
- with pytest.raises(RuntimeError):
- stackedtd.device
+ assert stackedtd.device is None
+
  stackedtd.set("a", torch.randn(2, 5, device=device))
- with pytest.raises(RuntimeError):
- stackedtd.device
+ assert stackedtd.device is None
 
  stackedtd = stackedtd.to(device)
  assert stackedtd.device == device
@@ -2139,12 +2134,12 @@ def test_create_on_device():
  # TensorDict, indexed
  td = TensorDict({}, [5])
  subtd = td[1]
- with pytest.raises(RuntimeError):
- subtd.device
+ assert subtd.device is None
+
  subtd.set("a", torch.randn(1, device=device))
- with pytest.raises(RuntimeError):
-  # setting element of subtensordict doesn't set top-level device
- subtd.device
+ # setting element of subtensordict doesn't set top-level device
+ assert subtd.device is None
+
  subtd = subtd.to(device)
  assert subtd.device == device
  assert subtd["a"].device == device
@@ -2162,8 +2157,8 @@ def test_create_on_device():
  # SavedTensorDict
  td = TensorDict({}, [5])
  savedtd = td.to(SavedTensorDict)
- with pytest.raises(RuntimeError):
- savedtd.device
+ assert savedtd.device is None
+
  savedtd = savedtd.to(device)
  assert savedtd.device == device
 
@@ -2175,8 +2170,8 @@ def test_create_on_device():
  # ViewedTensorDict
  td = TensorDict({}, [6])
  viewedtd = td.view(2, 3)
- with pytest.raises(RuntimeError):
- viewedtd.device
+ assert viewedtd.device is None
+
  viewedtd = viewedtd.to(device)
  assert viewedtd.device == device
 

torchrl/collectors/utils.py

@@ -41,7 +41,7 @@ def split_trajectories(rollout_tensordict: TensorDictBase) -> TensorDictBase:
  "mask",
  torch.ones(
  rollout_tensordict.shape,
- device=rollout_tensordict.device_safe(),
+ device=rollout_tensordict.device,
  dtype=torch.bool,
  ),
  )
@@ -66,7 +66,7 @@ def split_trajectories(rollout_tensordict: TensorDictBase) -> TensorDictBase:
  }
  td = TensorDict(
  source=out_dict,
- device=rollout_tensordict.device_safe(),
+ device=rollout_tensordict.device,
  batch_size=out_dict["mask"].shape[:-1],
  )
  if (out_dict["done"].sum(1) > 1).any():

torchrl/data/replay_buffers/replay_buffers.py

@@ -75,10 +75,7 @@ def stack_tensors(list_of_tensor_iterators: List) -> Tuple[torch.Tensor]:
 
 
 def _pin_memory(output: Any) -> Any:
- output_device = (
- output.device_safe() if hasattr(output, "device_safe") else output.device
- )
- if hasattr(output, "pin_memory") and output_device == torch.device("cpu"):
+ if hasattr(output, "pin_memory") and output.device == torch.device("cpu"):
  return output.pin_memory()
  else:
  return output
@@ -369,11 +366,7 @@ def __getitem__(self, index: Union[int, Tensor]) -> Any:
  weight = np.power(weight / p_min, -self._beta)
  # x = first_field(data)
  # if isinstance(x, torch.Tensor):
- device = (
- data.device_safe()
- if hasattr(data, "device_safe")
- else (data.device if hasattr(data, "device") else torch.device("cpu"))
- )
+ device = data.device if hasattr(data, "device") else torch.device("cpu")
  weight = to_torch(weight, device, self._pin_memory)
  return data, weight
 
@@ -476,11 +469,7 @@ def _sample(self, batch_size: int) -> Tuple[Any, torch.Tensor, torch.Tensor]:
 
  # x = first_field(data) # avoid calling tree.flatten
  # if isinstance(x, torch.Tensor):
- device = (
- data.device_safe()
- if hasattr(data, "device_safe")
- else (data.device if hasattr(data, "device") else torch.device("cpu"))
- )
+ device = data.device if hasattr(data, "device") else torch.device("cpu")
  weight = to_torch(weight, device, self._pin_memory)
  return data, weight, index
 
@@ -672,7 +661,7 @@ def extend(
  "index",
  torch.zeros(
  tensordicts.shape,
- device=tensordicts.device_safe(),
+ device=tensordicts.device,
  dtype=torch.int,
  ),
  )
@@ -686,7 +675,7 @@ def extend(
  idx = super().extend(tensordicts, priorities)
  stacked_td.set(
  "index",
- torch.tensor(idx, dtype=torch.int, device=stacked_td.device_safe()),
+ torch.tensor(idx, dtype=torch.int, device=stacked_td.device),
  inplace=True,
  )
  return idx

torchrl/data/tensordict/metatensor.py

@@ -81,15 +81,7 @@ def __init__(
  _is_memmap = isinstance(tensor, MemmapTensor)
  # FIXME: using isinstance(tensor, TensorDictBase) would likely be
  # better here, but creates circular import without more refactoring
- device = (
- (
- tensor.device_safe()
- if hasattr(tensor, "device_safe")
- else tensor.device
- )
- if not tensor.is_meta
- else device
- )
+ device = tensor.device if not tensor.is_meta else device
  if _is_tensordict is None:
  _is_tensordict = not _is_memmap and not isinstance(tensor, torch.Tensor)
  if not _is_tensordict: