Fix dictionary iteration

flferretti · Feb 2, 2024 · 7bfacee · 7bfacee
1 parent f5d189e
commit 7bfacee
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/jaxsim/simulation/simulator.py b/src/jaxsim/simulation/simulator.py
@@ -213,7 +213,7 @@ def get_model(self, model_name: str) -> Model:
             The model with the given name.
         """
 
-        if model_name not in self.data.models.keys():
+        if model_name not in self.data.models:
             raise ValueError(f"Failed to find model '{model_name}'")
 
         return self.data.models[model_name]
@@ -250,7 +250,7 @@ def set_gravity(self, gravity: jtp.Vector) -> None:
 
         self.data.gravity = gravity
 
-        for model_name, model in self.data.models.items():
+        for model in self.data.models.values():
             model.physics_model.set_gravity(gravity=gravity)
 
     @functools.partial(oop.jax_tf.method_rw, jit=False, vmap=False, validate=False)