formatting

mlcommons · Nov 18, 2023 · f39018c · f39018c
1 parent 49e221b
commit f39018c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/algorithmic_efficiency/workloads/mnist/mnist_pytorch/workload.py b/algorithmic_efficiency/workloads/mnist/mnist_pytorch/workload.py
@@ -80,9 +80,7 @@ def _build_input_queue(
           weights = torch.as_tensor(
               batch['weights'], dtype=torch.bool, device=DEVICE)
         else:
-          weights = torch.ones_like(targets, 
-                                    dtype=torch.bool, 
-                                    device=DEVICE)
+          weights = torch.ones_like(targets, dtype=torch.bool, device=DEVICE)
         # Send batch to other devices when using DDP.
         if USE_PYTORCH_DDP:
           dist.broadcast(inputs, src=0)

diff --git a/algorithmic_efficiency/workloads/mnist/workload.py b/algorithmic_efficiency/workloads/mnist/workload.py
@@ -214,8 +214,6 @@ def _eval_model_on_split(self,
                                        batch,
                                        model_state,
                                        per_device_model_rngs)
-      total_metrics = {
-          k: v + batch_metrics[k] for k, v in total_metrics.items()
-      }
+      total_metrics = {k: v + batch_metrics[k] for k, v in total_metrics.items()}
 
     return self._normalize_eval_metrics(num_examples, total_metrics)