Fixed #35

Spijkervet · Jul 24, 2021 · 04bcf2b · 04bcf2b
1 parent 5b4bdc8
commit 04bcf2b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/linear_evaluation.py b/linear_evaluation.py
@@ -195,13 +195,13 @@ def test(args, loader, simclr_model, model, criterion, optimizer):
             args, arr_train_loader, simclr_model, model, criterion, optimizer
         )
         print(
-            f"Epoch [{epoch}/{args.logistic_epochs}]\t Loss: {loss_epoch / len(train_loader)}\t Accuracy: {accuracy_epoch / len(train_loader)}"
+            f"Epoch [{epoch}/{args.logistic_epochs}]\t Loss: {loss_epoch / len(arr_train_loader)}\t Accuracy: {accuracy_epoch / len(arr_train_loader)}"
         )
 
     # final testing
     loss_epoch, accuracy_epoch = test(
         args, arr_test_loader, simclr_model, model, criterion, optimizer
     )
     print(
-        f"[FINAL]\t Loss: {loss_epoch / len(test_loader)}\t Accuracy: {accuracy_epoch / len(test_loader)}"
+        f"[FINAL]\t Loss: {loss_epoch / len(arr_test_loader)}\t Accuracy: {accuracy_epoch / len(arr_test_loader)}"
     )