diff --git a/src/jaxsim/simulation/simulator.py b/src/jaxsim/simulation/simulator.py
index 8c9e20249..77b370205 100644
--- 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)