Hi ptrblck,
Thank you very much for your reply. I have never used the cpp API and I am not sure how the python Module gets updated with the new value of current_mean but for writing a custom Python implementation, I guess the only thing I need to know (for now) is that I have to do it manually. Thanks for the tip.
Now for the practical part; if I want to keep it as close to the original structure as possible, and also as done in the example of the “extending Pytorch tutorial”, I have to create an autograd Function what will do the calculation and an nn Module that will call this function in forward(). However, since the current_mean would be calculated inside the Function (as done in the original batchnorm structure), how to allow the Module to keep track of it? Should I simply calculate the current_mean inside the Module?
I am sorry if I am not very clear, it is my first customizing experience in pyTorch, I am a physicist attempting a conversion…
Thanks again for your help !
SMP