[BugFix] Fix index in list error (#627)

pytorch · Jan 17, 2024 · 4225911 · 4225911
1 parent e50f614
commit 4225911
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -1838,7 +1838,8 @@ def format_size(size):
 def _index_preserve_data_ptr(index):
     if isinstance(index, tuple):
         return all(_index_preserve_data_ptr(idx) for idx in index)
-    if index in (None, Ellipsis):
+    # we can't use a list comprehension here because it fails with tensor indices
+    if index is None or index is Ellipsis:
         return True
     if isinstance(index, int):
         return True