Summary

• The inner error() function in classification_error was referencing probabilities from the enclosing scope instead of using its own predictions parameter
• This meant torch.argmax was applied to the closure variable rather than the argument passed by compute_masked_loss, which is incorrect — the function should operate on its own parameter
• Changed torch.argmax(probabilities, dim=-1) to torch.argmax(predictions, dim=-1)

Details

compute_masked_loss calls loss_fn(predictions, targets) and passes the predictions tensor as the first argument. The error function receives this as its predictions parameter, but then ignores it and captures probabilities from the outer scope via closure. While in the current call site these happen to be the same object, this is a latent bug:

1. The function parameter is shadowed/unused, violating the function contract
2. If compute_masked_loss is ever updated to transform predictions before calling the loss function, this would silently produce wrong results
3. Static analysis tools flag this as a bug (unused parameter + unintended closure capture)

Test plan

• Existing doctest classification_error(probs, torch.tensor([1, 1])) produces tensor(0.5000) (unchanged)
• Verify with 2D and 3D inputs that classification error is computed correctly