If I understand correctly if you put the with torch.no_grad() line the out value is the one that will be used to compute the gradients.  Since the operation out * out will not be added to the computational graph.  Again I am not 100% sure about this. Maybe @ptrblck can offer his alwaysvaluable insight. Best of luck.