Hi,
Your custom loss only needs to compute the partial derivative wrt it’s input. The autograd engine will use the chain rule and the other gradients formulas implemented in pytorch to compute the gradienst wrt the weights.
What you are given, if your loss is L and you compute out = your_fn(inp), is grad_output is dL/dout. And you should return dL/dinp. That is usually written as dL/dinp = dL/dout * dout/dinp. which is the formula you present above.
The autograd will take care of computing dL/dw using your computed dL/dinp.