If you loss function is differentiable, you don’t actually need to implement a Function. Just write an nn.Module that contains the forward pass and the gradients will be computed for you  (if they exist).
If it’s not differentiable, or the gradients that you want are not the one for the loss you compute, then yes creating a Function that way is the way to go.