amend

torchrl/data/tensor_specs.py

@@ -445,6 +445,7 @@ def __eq__(self, other):
  type(self) == type(other)
  and self.low.dtype == other.low.dtype
  and self.high.dtype == other.high.dtype
+ and self.device == other.device
  and torch.isclose(self.low, other.low).all()
  and torch.isclose(self.high, other.high).all()
  )
@@ -1004,6 +1005,8 @@ class LazyStackedTensorSpec(_LazyStackedMixin[TensorSpec], TensorSpec):
  def __eq__(self, other):
  if not isinstance(other, LazyStackedTensorSpec):
  return False
+ if self.device != other.device:
+ return False
  if len(self._specs) != len(other._specs):
  return False
  for _spec1, _spec2 in zip(self._specs, other._specs):
@@ -1593,6 +1596,7 @@ def __init__(
  def __eq__(self, other):
  return (
  type(other) == type(self)
+ and self.device == other.device
  and self.shape == other.shape
  and self.space == other.space
  and self.dtype == other.dtype
@@ -2825,9 +2829,9 @@ def __eq__(self, other):
  if isinstance(other, DiscreteTensorSpec):
  return (
  other.n == 2
+ and other.device == self.device
  and other.shape == self.shape
  and other.dtype == self.dtype
- and other.device == self.device
  )
  return False
  return super().__eq__(other)
@@ -3715,7 +3719,7 @@ def __eq__(self, other):
  return (
  type(self) is type(other)
  and self.shape == other.shape
- and self._device == other._device
+ and self.device == other.device
  and set(self._specs.keys()) == set(other._specs.keys())
  and all((self._specs[key] == spec) for (key, spec) in other._specs.items())
  )
@@ -3937,6 +3941,8 @@ def __eq__(self, other):
  return False
  if self.stack_dim != other.stack_dim:
  return False
+ if self.device != other.device:
+ return False
  for _spec1, _spec2 in zip(self._specs, other._specs):
  if _spec1 != _spec2:
  return False