Update _lazy.py

pytorch · Feb 3, 2024 · e722059 · e722059
1 parent 5302ac1
commit e722059
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
@@ -2040,9 +2040,7 @@ def update_(
             # no op
             return self
         if not is_tensor_collection(input_dict_or_td):
-            input_dict_or_td = TensorDict.from_dict(
-                input_dict_or_td, batch_dims=self.batch_dims
-            )
+            input_dict_or_td = TensorDict.from_dict(input_dict_or_td)
         if input_dict_or_td.batch_size[self.stack_dim] != len(self.tensordicts):
             raise ValueError("cannot update stacked tensordicts with different shapes.")
         for td_dest, td_source in zip(