Lint fix

mlcommons · Oct 2, 2023 · a1844c7 · a1844c7
1 parent 19accc7
commit a1844c7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/algorithmic_efficiency/workloads/criteo1tb/criteo1tb_jax/workload.py b/algorithmic_efficiency/workloads/criteo1tb/criteo1tb_jax/workload.py
@@ -148,7 +148,8 @@ def _eval_batch(self,
                   batch: Dict[str, spec.Tensor]) -> spec.Tensor:
     # We do NOT psum inside of _eval_batch_pmapped, so the returned tensor of
     # shape (local_device_count,) will all be different values.
-    return np.array(self._eval_batch_pmapped(params, batch).sum(), dtype=np.float64)
+    return np.array(
+        self._eval_batch_pmapped(params, batch).sum(), dtype=np.float64)
 
 
 class Criteo1TbDlrmSmallTestWorkload(Criteo1TbDlrmSmallWorkload):