[Optimization] Replace torch.no_grad() with torch.inference_mode context

[fedebotu]

torch.inference_mode is a new context manager that seems to be faster than torch.no_grad(), so we may replace calls to that

[fedebotu]

Released in v0.3.2 - noticed something like 1% speedup per epoch, which is not bad at all for such simple change 💪