_zip_strict

pytorch · Oct 9, 2024 · 8dc35ec · 8dc35ec
1 parent 8093908
commit 8dc35ec
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -22,6 +22,7 @@
 import torch
 
 from tensordict import (
+    _zip_strict,
     is_tensor_collection,
     LazyStackedTensorDict,
     TensorDict,
@@ -544,9 +545,7 @@ def update_kwargs(self, kwargs: Union[dict, List[dict]]) -> None:
                 raise RuntimeError(
                     f"len(kwargs) and num_workers mismatch, got {len(kwargs)} and {self.num_workers}."
                 )
-            for _kwargs, _new_kwargs in zip(
-                self.create_env_kwargs, kwargs, strict=True
-            ):
+            for _kwargs, _new_kwargs in _zip_strict(self.create_env_kwargs, kwargs):
                 _kwargs.update(_new_kwargs)
 
     def _get_in_keys_to_exclude(self, tensordict):