[BugFix] Wrong call to device_safe in replay buffer code (#454)

Author: vmoens

torchrl/data/replay_buffers/replay_buffers.py

@@ -75,7 +75,10 @@ def stack_tensors(list_of_tensor_iterators: List) -> Tuple[torch.Tensor]:
 
 
 def _pin_memory(output: Any) -> Any:
-    if hasattr(output, "pin_memory") and output.device_safe() == torch.device("cpu"):
+    output_device = (
+        output.device_safe() if hasattr(output, "device_safe") else output.device
+    )
+    if hasattr(output, "pin_memory") and output_device == torch.device("cpu"):
         return output.pin_memory()
     else:
         return output