Solved by Sagnik_Mukherjee in post #6 
              
              
                Okay, I can help you a bit there. Gradient is a “fancy” name for derivative (differential calculus). so lets say you have a tensor (2*2) defined as: 
W11 W12 
W21 W22 
When you take torch.mean(), the returned variable say W = (W11 + W21 + W12 + W22)/4.0 
And hence the gradient matrix shall look like…