Thanks a lot!
I am trying to implement that but I am stuck in another problem.
My model is quite deep so I have multiple batchnorm layers, I thought about creating copies of them and save in a ModuleDict.
But on forward I need to go through all batchnorm layers and assign them the desired layers. This I dont know how to. I thought it would be nice to access the layers by name, something like:
model.modules[name] = synth_batchnorm.
I tried also something like model.named_modules()[i] = synth_batchnorm, but I get the error ‘generator object does not support item assignment’.
How that can be done?