Thanks. My issue is that I have a really annoying use case (in meta learning) where this hack doesn’t work. All my weights must be concatenated into a single trainable tensor, and passed to the network as an argument when doing the forward pass. So this trick doesn’t work because that would mean having some elements in my weight vector be non-trainable.
In short I need gammas not to be parameters when affine=False otherwise they break everything… Do you think tweaking the BN layer source code is required here? I can’t figure out why gamma=None gets interpreted as gamma=1 in the BN function.