Appease type checker

JasonGross · Sep 5, 2024 · fc5bf4a · fc5bf4a
1 parent 23a0a19
commit fc5bf4a
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/notebooks_jason/max_of_K_all_models.py b/notebooks_jason/max_of_K_all_models.py
@@ -679,8 +679,7 @@ def _run_train_batch_loss_accuracy(
     loss, accuracy = training_wrappers[seed].run_batch(
         (xs, ys), log_output=False, device=device
     )
-    loss = loss.item()
-    return loss, accuracy, batch_size
+    return loss.item(), accuracy, batch_size
 
 
 def train_seed(seed: int, *, memoshelve_hf: Callable, pbar: tqdm):