Update.

diff --git a/main.py b/main.py
index 63819f2..6a7ca3f 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -323,7 +323,7 @@ def run_tests(model, quiz_machine, deterministic_synthesis, local_device=main_de
 
         test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples))
 
-        log_string(f"test_perplexity {n_epoch} {test_perplexity}")
+        log_string(f"test_perplexity {n_epoch} model {model.id} {test_perplexity}")
 
         model.main_test_accuracy = quiz_machine.produce_results(
             n_epoch=n_epoch,
@@ -359,10 +359,12 @@ def one_epoch(model, quiz_machine, local_device=main_device):
 
     train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples))
 
-    log_string(f"train_perplexity {n_epoch} model.id {model.id} {train_perplexity}")
+    log_string(f"train_perplexity {n_epoch} model {model.id} {train_perplexity}")
 
     run_tests(model, quiz_machine, deterministic_synthesis=False)
 
+    model.to(main_device)
+
 
 ######################################################################