I also spotted that gradients are not accumulated correctly, so I fixed that too.

Thank you for this PR @MartinKocour. I think it can be useful. One thing I'm not sure about is why calling the function you defined (zero_grad) in fit_batch as well. As you can see, in core.py you still call the original function self.optimizer.zero_grad(), but you modified the recipes files such that they call zero_grad in their fit_batch . I'm not sure we need to delete the gradient buffers at the end of each batch (this might slow down?), but likely we really need to do it at the end of a training epoch only, right?

Mirco @mravanelli, in recipes I am not deleting anything. My zero_grad(set_to_none=False) has the same behaviour like original torch.optimizer.zero_grad(), i.e. by default it does not delete the gradients from memory, but instead it just zeros the gradients when set_to_none=False (default).

You are right, I don't have to call self.zero_grad() in the recipes. I just didn't want to duplicate same lines of code. But it is not necessary, I can remove it. Although, the behaviour with self.zero_grad() is the same like with original code.

In core.py I did not use it since it is just single line of code + calling optimizer.zero_grad() is faster than calling Brain.zero_grad() since it contains additional if statement.

Ready to be merged