fix pos output

sudoku/train.py

@@ -875,7 +875,7 @@ class SudokuTrialErrorLightning(SudokuLightning):
             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 = torch.argmax(output[0,:,pos])
+            pos_output = output[0,:,pos]
             if (pos_output[0].item()-self.threshold_abs[0].item())>(pos_output[1].item-self.threshold_pos[0].item()):
                 trial_abs_pres = 0
             else: