Fix softmax masking for illegal moves in inference

Multiplying logits by a binary 0/1 mask sets illegal move logits to 0,
but softmax(0) is not zero — it still contributes probability mass.
Use masked_fill with -inf so softmax correctly assigns exactly 0
probability to illegal moves.

Fixes the issue raised in PR #9.

https://claude.ai/code/session_01ELpknikQ4vWB4q8hWhxoVG

Author: claude

maia2/inference.py

@@ -60,7 +60,7 @@ def get_preds(model, dataloader, all_moves_dict_reversed):
             legal_moves = legal_moves.to(device)
 
             logits_maia, _, logits_value = model(boards, elos_self, elos_oppo)
-            logits_maia_legal = logits_maia * legal_moves
+            logits_maia_legal = logits_maia.masked_fill(legal_moves == 0, float('-inf'))
             probs = logits_maia_legal.softmax(dim=-1).cpu().tolist()
 
             logits_value = (logits_value / 2 + 0.5).clamp(0, 1).cpu().tolist()
@@ -154,7 +154,7 @@ def inference_each(model, prepared, fen, elo_self, elo_oppo):
     legal_moves = legal_moves.unsqueeze(dim=0).to(device)
 
     logits_maia, _, logits_value = model(board_input, elo_self, elo_oppo)
-    logits_maia_legal = logits_maia * legal_moves
+    logits_maia_legal = logits_maia.masked_fill(legal_moves == 0, float('-inf'))
     probs = logits_maia_legal.softmax(dim=-1).cpu().tolist()
 
     logits_value = (logits_value / 2 + 0.5).clamp(0, 1).item()