Hdk:

So, I have to do create weight,bias, running_mean and running_var for each layer in init , right?
In your code, I am guessing there’s no storing is done, can you please inform where exactly and how to store?


You can store these parameters and buffers in the __init__ method of an nn.Module, but then again you wouldn’t use the functional API but modules.
If you want to store these tensors in a module why wouldn’t you want to use the module (nn.Conv2d, nn.BatchNorm2d) directly?