You can define your own loss function by extending nn.Module class.

import torch.nn as nn
import torch


class Loss(nn.Module):
    def __init__(self):
        
        super(Loss, self).__init__()
        

    def forward(self, y, y_pred):
        loss = # DO YOUR CALCULATION ON THE TENSOR OF YOUR IMAGE
        return loss