You can use hooks. So for linear layers, you could do this:

  
      gist.github.com
  
  
    https://gist.github.com/t-vi/f3437d31b3e4680cc78d9999ea5a8af6#file-variance_of_grad-py 57
variance_of_grad.py
import torch
from torch.autograd import Variable

def linear_with_sumsq(inp, weight, bias=None):
    def provide_sumsq(inp,w,b):
        def _h(i):
            if not hasattr(w, 'grad_sumsq'):
                w.grad_sumsq = 0
            w.grad_sumsq += ((i**2).t().matmul(inp**2))*i.size(0)
            if b is not None:
This file has been truncated. show original




  
  
    
    
  
  


For convolutions, you would need to instantiate an object of the ConvNdBackwardBackward class, which I don’t think is possible at the moment. (But might be in the future if standard deviations become more popular.)
Best regards
Thomas