I’m trying to implement a “BatchScale” Module that is akin to BatchNorm, except that there is no shift/mean/bias, just a scaling/weighting and a running variance.  Initially I thought I could just copy and modify the _BatchNorm class, but it calls the C version of batch_norm and as far as I can tell, that function will always compute and apply a running mean.
Any suggestions?