Solved by kgbounce in post #10 
              
              
                Hey, I figured out how to do it with autograd.Function after reading through the pages. Let me know if it makes sense. 
class TwoOpGradientController(torch.autograd.Function):
    @staticmethod
    def forward(ctx, *args):
        ctx.args_0_shape = args[0].shape
        return args[0] @ args[1]
   …