Update torchrl/data/replay_buffers/replay_buffers.py

Co-authored-by: Vincent Moens <[email protected]>
pytorch · Oct 3, 2023 · bb74238 · bb74238
1 parent 81dcee7
commit bb74238
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/data/replay_buffers/replay_buffers.py b/torchrl/data/replay_buffers/replay_buffers.py
@@ -694,7 +694,7 @@ def _get_priority_vector(self, tensordict: TensorDictBase) -> torch.Tensor:
                 device=tensordict.device,
             ).expand(tensordict.shape[0])
 
-        priority = priority.view(priority.shape[0], -1)
+        priority = priority.reshape(priority.shape[0], -1)
         priority = _reduce(priority, self._sampler.reduction, dim=1)
 
         return priority