diff --git a/deep_quoridor/src/agents/alphazero/mlp_network.py b/deep_quoridor/src/agents/alphazero/mlp_network.py
index 3fc83cb5..f2f94a00 100644
--- a/deep_quoridor/src/agents/alphazero/mlp_network.py
+++ b/deep_quoridor/src/agents/alphazero/mlp_network.py
@@ -24,7 +24,9 @@ def __init__(self, input_size, action_size, device):
         # TODO: Is it correct to include the Softmax at the end? Some implementations of alphazero
         # appear to leave it out, or apply it outside the network.
         self.policy_head = nn.Sequential(
-            nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, action_size), nn.Softmax(dim=-1)
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Linear(128, action_size),
         )
 
         # value head - outputs position evaluation (-1 to 1)
diff --git a/deep_quoridor/src/agents/alphazero/nn_evaluator.py b/deep_quoridor/src/agents/alphazero/nn_evaluator.py
index 31ef558e..74ae178e 100644
--- a/deep_quoridor/src/agents/alphazero/nn_evaluator.py
+++ b/deep_quoridor/src/agents/alphazero/nn_evaluator.py
@@ -36,26 +36,21 @@ def evaluate(self, game: Quoridor):
 
         self.network.eval()  # Disables dropout
 
+        valid_actions = game.get_valid_actions()
+        valid_action_indices = [self.action_encoder.action_to_index(action) for action in valid_actions]
+
         with torch.no_grad():
             input_array = self.game_to_input_array(game)
-            unmasked_policy, value = self.network(torch.from_numpy(input_array).float().to(self.device))
-            unmasked_policy = unmasked_policy.cpu().numpy()
-            value = value.item()
+            unmasked_policy_logits, value = self.network(torch.from_numpy(input_array).float().to(self.device))
+            unmasked_policy_logits = unmasked_policy_logits.cpu().numpy()
 
-        # Mask the policy to ignore invalid actions. NOTE: Game is already rotated so the valid actions will be rotated too
-        valid_actions = game.get_valid_actions()
-        valid_action_indices = [self.action_encoder.action_to_index(action) for action in valid_actions]
-        policy_masked = np.zeros_like(unmasked_policy)
-        policy_masked[valid_action_indices] = unmasked_policy[valid_action_indices]
-
-        if np.all(policy_masked == 0):
-            # If the policy ends up as all zeros after masking, turn it into a uniform distribution among
-            # the valid actions.
-            policy_masked[valid_action_indices] = 1 / len(valid_action_indices)
-            print("Policy is all zeros after masking, turning it into a uniform distribution")
-        else:
-            # Otherwise, just renormalize after masking
-            policy_masked = policy_masked / policy_masked.sum()
+            # Set logits of invalid actions to negative infinity to ensure they are never chosen
+            # TODO: Mask entirely in Torch to avoid going to Numpy and back
+            masked_logits = np.full_like(unmasked_policy_logits, -np.inf)
+            masked_logits[valid_action_indices] = unmasked_policy_logits[valid_action_indices]
+
+            policy_masked = F.softmax(torch.from_numpy(masked_logits), dim=-1).cpu().numpy()
+            value = value.item()
 
         # Sanity checks
         assert np.all(policy_masked >= 0), "Policy contains negative probabilities"
@@ -159,10 +154,11 @@ def train_iteration(self, replay_buffer):
                 raise ValueError("NaN in training data")
 
             # Forward pass
-            pred_policies, pred_values = self.network(inputs)
+            pred_logits, pred_values = self.network(inputs)
+            # TODO: Should we apply masking before calculating cross-entropy here?
 
             # Compute losses
-            policy_loss = F.cross_entropy(pred_policies, target_policies, reduction="mean")
+            policy_loss = F.cross_entropy(pred_logits, target_policies, reduction="mean")
             value_loss = F.mse_loss(pred_values.squeeze(), target_values.squeeze(), reduction="mean")
             total_loss = policy_loss + value_loss
             # print(f"{total_loss.item():3.3f} {policy_loss.item():3.3f} {value_loss.item():3.3f}")