fixed pmap call

AshishKumar4 · Aug 4, 2024 · 25c7c16 · 25c7c16
1 parent 6b87ce1
commit 25c7c16
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/training.py b/training.py
@@ -540,7 +540,7 @@ def model_loss(params):
             return train_state, loss, rng_state
 
         if distributed_training:
-            train_step = jax.pmap(axis_name="data")(train_step)
+            train_step = jax.pmap(train_step, axis_name="data")
             # train_step = shard_map(train_step, mesh=self.mesh, in_specs=P('data'), out_specs=P())
         else:
             train_step = jax.jit(train_step)
@@ -811,7 +811,7 @@ def model_loss(params):
             return train_state, loss, rng_state
 
         if distributed_training:
-            train_step = jax.pmap(axis_name="data")(train_step)
+            train_step = jax.pmap(train_step, axis_name="data")
             # train_step = shard_map(train_step, mesh=self.mesh, in_specs=P('data'), out_specs=P())
         else:
             train_step = jax.jit(train_step)