[BugFix]: Tensor map for subtensordict.set_ (#324)

Author: vmoens

torchrl/data/replay_buffers/replay_buffers.py

@@ -670,7 +670,11 @@ def extend(
         else:
             stacked_td = tensordicts
         idx = super().extend(tensordicts, priorities)
-        stacked_td.set("index", idx, inplace=True)
+        stacked_td.set(
+            "index",
+            torch.tensor(idx, dtype=torch.int, device=stacked_td.device),
+            inplace=True,
+        )
         return idx
 
     def update_priority(self, tensordict: TensorDictBase) -> None:

torchrl/data/tensordict/tensordict.py

@@ -2766,6 +2766,9 @@ def set_(
                     f"tensor.shape={tensor.shape[:self.batch_dims]} and "
                     f"self.batch_size={self.batch_size} mismatch"
                 )
+        tensor = self._process_tensor(
+            tensor, check_device=False, check_tensor_shape=False
+        )
         self._source.set_at_(key, tensor, self.idx)
         if key in self._dict_meta:
             self._dict_meta[key].requires_grad = tensor.requires_grad