Update.

[culture.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index cb5900b..9a67127 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -112,6 +112,13 @@ class World(Task):
             nb_test_samples, height=self.height, width=self.width
         ).to(device)
 
+        # print()
+        # for a in world.seq2str(self.train_input):
+        # print(a)
+        # for a in world.seq2str(self.test_input):
+        # print(a)
+        # exit(0)
+
         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 
         self.train_quizzes = []
@@ -274,7 +281,7 @@ class World(Task):
         # Check how many of the other models can solve them in both
         # directions
 
-        nb_correct = 0
+        nb_correct = []
 
         for m in other_models:
             result = quizzes.clone()
@@ -307,6 +314,13 @@ class World(Task):
                 (reverse_quizzes == reverse_result).long().min(dim=-1).values
             )
 
-            nb_correct += correct * reverse_correct
+            nb_correct.append((correct * reverse_correct)[None, :])
+
+        nb_correct = torch.cat(nb_correct, dim=0)
+
+        filename = os.path.join(result_dir, "correct_{n_epoch:04d}.dat")
+        with open(filename, "w") as f:
+            for k in nb_correct:
+                f.write(f"{k}\n")
 
-        return quizzes, nb_correct
+        return quizzes, nb_correct.sum(dim=0)