Solved by albanD in post #2 
              
              
                Hi, 
You can use triple backticks ``` before and after your code to have nicer formatting. 
The combination of gradients if not an element wise product but a matrix matrix multiplications. 
You can do torch.bmm(dlogits_dnew_relu_feat, dloss_dlogits.unsqueeze(-1)).squeeze(-1) to get what you want. 
T…