Thanks a lot! Where should I exactly use create_graph=True? Do I really need it while I am not using autograd.grad?
Here is my code, where it is defining some new custom loss class with all the computations inside LossAndGradients:
class CustomLoss(Function): # Inherit from Function
    
    @staticmethod
    def forward(ctx, X_input, imgD):        
        Loss, Gradients = LossAndGradients(X_input, imgD)
        ctx.save_for_backward(Gradients)
        return Loss
    
    @staticmethod
    def backward(ctx, grad_output):
        Gradients, = ctx.saved_tensors
        grad_imgD = None
        grad_input = Gradients
        return grad_input, grad_imgD

In the above function I have set  requires_grad=True for Loss and X_input, and requires_grad=False for Gradients. And I employ it as below in the main code:
criterion = CustomLoss.apply

I am not really sure if it is back-propagating the gradients in a correct way during training.