Hi @albanD, sorry for my delayed answer. I hope you help me. I modified my function in many ways (using Function, using Module,…), but the final computed loss is only a tensor, without any grad_fn  (That is, it is out the computation graph).
My goal is get dice_coeff and backward it (the formula isn’t exact).
Can you allow me to show line by line?

criterion = DiceCoeffLoss() # custom loss function declaration, its code is below


outputs = model(inputs)

and his gradient function:
outputs.grad_fn Out[285]: <SigmoidBackward at 0x1b936013e08 # rigth

loss = criterion(outputs,targets)

outputs shape is [n_batch, n_class, H, W]
target shape is [n_batch,  H, W]
Into criterion, in order to calculate the loss, first I have to get the predicted labels along the channels:

pred = torch.max(outputs, dim=1)

Here, pred has no direct attribute grad_fn,  pred.values  yet contains the  outputs  values and pred.values.grad_fn has the gradient function  <MaxBackward0 at 0x1b93601d1c8>.
However, since next operations is based on the predicted labels contained in  pred.indices , the final  loss  is only a tensor.
Here the code (disclaimer, the formula isn’t the rigth):
def dice_coeff(predicted, targets):
    smooth = 1.0        
    inter = predicted.eq(targets).sum() # only true predited
    union = float(predicted.nelement() + targets.nelement() ) # all elements        
    dc = (2*inter + smooth) / (union + smooth)        
    return  dc

# I used nn.Module thinking it will add those operations to the graph    
class DiceCoeffLoss(nn.Module):
    def __init__(self):
        super(DiceCoeffLoss, self).__init__()        
         
    def forward(self, outputs, targets):
        
        pred = torch.max(outputs, dim=1)
        coeff = dice_coeff(pred.indices,targets)    
        return 1 - coeff