add notebook and fix code

sudoku/train.py

@@ -869,29 +869,37 @@ class SudokuTrialErrorLightning(SudokuLightning):
         """
         return is_valid, new_x
         """
-        try:
-            next_x = self.predict(x)
-        except TrialEveryPosException:
-            pos = self.search_trial(x.view(2,729), [])
-            new_x = deepcopy(x.view(2,729))
-            output = self.forward_layer(x.view(1,2,729))
-            pos_output = output[0,:,pos]
-            if (pos_output[0].item()-self.threshold_abs[0].item())>(pos_output[1].item-self.threshold_pres[0].item()):
-                trial_abs_pres = 0
+        next_X = deepcopy(x)
+        sum_1 = next_X.sum().item()
+        while True:
+            idx, next_X = self.forward(next_X.view(1,2,729))
+            if not self.validate_grids(next_X)[0].item():
+                return False, None
+            sum_2 = next_X.sum().item()
+            if sum_1==sum_2:
+                break
             else:
-                trial_abs_pres = 1
-            new_x[0,trial_abs_pres,pos]=1
-            is_valid, new_x = self.backtracking_predict(new_x)
-            if not is_valid:
-                new_x = deepcopy(x.view(2,729))
-                new_x[0,(trial_abs_pres+1)%2,pos]=1
-                return self.backtracking_predict(new_x)
-            return True, new_x
-        except ValueError:
-            return False, None
-        if next_x.sum()==729:
-            return True, next_x
-        return self.backtracking_predict(next_x)
+                sum_1=sum_2
+        if next_X.sum()>=729:
+            return True, next_X
+        
+        pos = self.search_trial(next_X.view(2,729), [])
+        new_x = deepcopy(next_X.view(1,2,729))
+        output = self.forward_layer(next_X.view(1,2,729))
+        pos_output = output[0,:,pos]
+        if (pos_output[0].item()-self.threshold_abs[0].item())>(pos_output[1].item()-self.threshold_pres[0].item()):
+            trial_abs_pres = 0
+        else:
+            trial_abs_pres = 1
+        new_x[0,trial_abs_pres,pos]=1
+        print(f'assuming {(pos, pos%9, (pos//9)%9, pos//(9*9), trial_abs_pres)}')# i+j*9+n*9*9
+        is_valid, new_x = self.backtracking_predict(new_x)
+        if not is_valid:
+            new_x = deepcopy(next_X.view(1,2,729))
+            new_x[0,(trial_abs_pres+1)%2,pos]=1
+            print(f'finally assuming {(pos, pos%9, (pos//9)%9, pos//(9*9), (trial_abs_pres+1)%2)}')# i+j*9+n*9*9
+            return self.backtracking_predict(new_x)
+        return True, new_x
             
     
     def on_validation_epoch_start(self) -> None: