X-Git-Url: https://www.fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=0858282fa3b5cab185d14025d5ad758de44411c2;hb=0f4c86c0e7730db4147f136df5aeb5528fc943a0;hp=b33dee2715621f4180d579afc3516416f51951ec;hpb=3e3bf1003aa0ecbf7d38b7b0c289fbe1cfa3101b;p=picoclvr.git

diff --git a/tasks.py b/tasks.py
index b33dee2..0858282 100755
--- a/tasks.py
+++ b/tasks.py
@@ -110,13 +110,14 @@ class SandBox(Task):
 
         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 
+
         # A bit of paranoia never hurts
         assert (
             self.nb_codes <= max_nb_codes
             and self.train_input.min() >= 0
             and self.test_input.min() >= 0
-            and tuple(self.train_ar_mask.unique()) == (0, 1)
-            and tuple(self.test_ar_mask.unique()) == (0, 1)
+            and tuple(x.item() for x in self.train_ar_mask.unique()) in { (0,), (1,), (0,1) }
+            and tuple(x.item() for x in self.test_ar_mask.unique()) in { (0,), (1,), (0,1) }
         )
 
     def batches(self, split="train", nb_to_use=-1, desc=None):
@@ -160,8 +161,10 @@ class SandBox(Task):
                         f"               {n_epoch} ground truth {self.problem.seq2str(st)}"
                     )
 
-            nb_total = ar_mask.sum().item()
-            nb_correct = ((result == input).long() * ar_mask).sum().item()
+            nb_total, nb_correct = self.problem.compute_nb_correct(input, ar_mask, result)
+
+            # nb_total = ar_mask.sum().item()
+            # nb_correct = ((result == input).long() * ar_mask).sum().item()
 
             return nb_total, nb_correct