[BugFix] Fix ellipsis check (#604)

pytorch · Jan 3, 2024 · 1b6b30c · 1b6b30c
1 parent 5b6573a
commit 1b6b30c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tensordict/_td.py b/tensordict/_td.py
@@ -526,7 +526,10 @@ def __setitem__(
                 )
                 return
 
-        if index is Ellipsis or (isinstance(index, tuple) and Ellipsis in index):
+        # we must use any and because using Ellipsis in index can break with some indices
+        if index is Ellipsis or (
+            isinstance(index, tuple) and any(idx is Ellipsis for idx in index)
+        ):
             index = convert_ellipsis_to_idx(index, self.batch_size)
 
         if isinstance(value, (TensorDictBase, dict)):