init

pytorch · Oct 10, 2023 · 2fb64be · 2fb64be
1 parent 7747505
commit 2fb64be
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
@@ -7924,6 +7924,8 @@ def sort_keys(element):
     rename_key = _renamed_inplace_method(rename_key_)
 
     def where(self, condition, other, *, out=None, pad=None):
+        if condition.ndim < self.ndim:
+            condition = expand_right(condition, self.batch_size)
         condition = condition.unbind(self.stack_dim)
         if _is_tensor_collection(other.__class__) or (
             isinstance(other, Tensor)