Update.

diff --git a/main.py b/main.py
index 0311450..0056c76 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -598,7 +598,9 @@ def one_complete_epoch(
 
     # Compute the test accuracy
 
-    quizzes = generate_quiz_set(args.nb_test_samples, c_quizzes, args.c_quiz_multiplier)
+    quizzes = generate_quiz_set(
+        args.nb_test_samples, test_c_quizzes, args.c_quiz_multiplier
+    )
     imt_set = samples_for_prediction_imt(quizzes.to(local_device))
     result = ae_predict(model, imt_set, local_device=local_device).to("cpu")
     correct = (quizzes == result).min(dim=1).values.long()
@@ -611,7 +613,11 @@ def one_complete_epoch(
     )
 
     save_inference_images(
-        model, n_epoch, c_quizzes, args.c_quiz_multiplier, local_device=local_device
+        model,
+        n_epoch,
+        train_c_quizzes,
+        args.c_quiz_multiplier,
+        local_device=local_device,
     )