small fix

priyakasimbeg · priyakasimbeg · commit 0a340a2f9f82 · 2025-04-03T22:29:17.000Z
diff --git a/algoperf/workloads/cifar/cifar_jax/workload.py b/algoperf/workloads/cifar/cifar_jax/workload.py
@@ -183,7 +183,7 @@ def _compute_metrics(self,
               jax_sharding_utils.get_replicated_sharding(),  # params
               jax_sharding_utils.get_batch_sharding(),  # batch
               jax_sharding_utils.get_replicated_sharding(),  # model_state
-              jax_sharding  _utils.get_batch_sharding(),  # rng
+              jax_sharding_utils.get_batch_sharding(),  # rng
           ),
       )
   def _eval_model(

Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,7 @@ def _compute_metrics(self,`
`183`	`183`	`jax_sharding_utils.get_replicated_sharding(), # params`
`184`	`184`	`jax_sharding_utils.get_batch_sharding(), # batch`
`185`	`185`	`jax_sharding_utils.get_replicated_sharding(), # model_state`
`186`		`- jax_sharding _utils.get_batch_sharding(), # rng`
	`186`	`+ jax_sharding_utils.get_batch_sharding(), # rng`
`187`	`187`	`),`
`188`	`188`	`)`
`189`	`189`	`def _eval_model(`