Fixes #1710.

Description

The problem was that for class_idx=None, the argmax was used in training_mode, resulting in different inferred classes compared to eval_mode. Secondly, the model was getting silently updated. Solved by using self.model.eval() but not using with torch.no_grad() or with eval_mode(self.model).

I've added a new test to check that when class_idx=None the inferred model is the same as that using eval_mode.

Status

Ready

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh --codeformat --coverage.
• Quick tests passed locally by running ./runtests.sh --quick.