yea, we could try caching them for inference. that doesnt need to be done in the first version of this pr but it might be a nice freezing speed up. the mean and variance are wrt the weights right ? edit: actually, might be wrt/ the input so nevermind..