thanks for clarifying. I am still not completely there.
so let’s say I have the following:
‘’’
class SDTWLoss(torch.autograd.Function):
@staticmethod
def forward( ctx, y_pred, y):
   
    _dtw = SoftDTW(y_pred, y ,device)
    dtw_value, dtw_grad, dtw_E_mat= _dtw.compute() # calculate loss value for y_prediction and y_actual
    G = jacobian_product_with_expected_sdtw_grad(y_pred, y, dtw_grad,device) # calculate derivative of loss w.r.t. y_pred (last layer of my graph) 
    ctx.save_for_backward(G) #save derviative of Loss w.r.t y for backward
    dtw_loss = torch.mean(dtw_value)

    return dtw_loss


@staticmethod
def backward( ctx, grad_output ):

    G = ctx.saved_tensors
    grad_input = grad_output.clone() # is this the derivative of the layer before the last layer w.r.t model's parameters???
    # if so, I should multiply G by grad_input here, right?
    grad_input = G*grad_input

    return grad_input

‘’’
Does this code make sense?