fixed test accuracy tuple

LostOxygen · Jul 30, 2023 · bfa60a3 · bfa60a3
1 parent f7d6925
commit bfa60a3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kernel_eval/train.py b/kernel_eval/train.py
@@ -128,7 +128,7 @@ def train_model(model: nn.Module, train_dataloader: IterableDataset,
         train_losses.append(sum(epoch_loss) / len(epoch_loss))
 
         # for every epoch use the validation set to check if this is the best model yet
-        validation_acc = test_model(model, validation_dataloader, device)
+        validation_acc = test_model(model, validation_dataloader, device)[0]
         wandb.log({"validation_acc": validation_acc})
 
         if validation_acc > best_validation_acc: