Merge remote-tracking branch 'origin/release/0.3.2' into release/0.3.2

pytorch · Apr 7, 2024 · 3ed6e8b · 3ed6e8b
2 parents 9301cb6 + 6a18431
commit 3ed6e8b
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/torchrl/data/replay_buffers/utils.py b/torchrl/data/replay_buffers/utils.py
@@ -10,6 +10,7 @@
 
 import numpy as np
 import torch
+
 from torch import Tensor
 
 INT_CLASSES_TYPING = Union[int, np.integer]
@@ -87,3 +88,11 @@ def _reduce(
     if isinstance(result, tuple):
         result = result[0]
     return result.item() if dim is None else result
+
+
+def _is_int(index):
+    if isinstance(index, INT_CLASSES):
+        return True
+    if isinstance(index, (np.ndarray, torch.Tensor)):
+        return index.ndim == 0
+    return False