How should I implement my very own custom loss function.
There are two ways to implement using a custom loss class and writing function with init and forward.
I had the code already written for TensorFlow in class, therefore I chose to go along.
For PyTorch, I tried explored the below relevant discussions, but could not figure how to design the backwards function for same.

Post_1 1
Post_2 1
Post_3 1

Custom Loss:
class MSE_ITL_RegLoss(torch.nn.Module):
    def __init__(self, latent_code, prior, regularization_factor, kernel, batch_size):
        super (MSE_ITL_RegLoss).__init__()
        
        self.latent_code = latent_code         
        self.regularization_factor = regularization_factor
        self.prior = prior
        self.kernel = kernel
        self.batch_size = batch_size
        
    
    def evaluate_ITL(self, first_pdf, second_pdf):       
        error = first_pdf - second_pdf
        pairwise_error = torch.reshape(torch.unsqueeze(error ,1) - error, (-1,))
        temp = torch.exp(-(np.sqrt(2) * self.kernel) * torch.square(pairwise_error_t))
        return torch.sum(temp)/self.batch_size**2
        
    
    def cs_div (self):
        p_x = self.evaluate_ITL(self.latent_code, self.latent_code)
        q_x = self.evaluate_ITL(self.prior,self.prior)
        p_x__q_x = self.evaluate_ITL(self.latent_code, self.prior) 
        return p_x__q_x / torch.sqrt(p_x * q_x)
        
        
    def forward(self, y_true, y_pred):       
        mse = torch.mean(torch.square(y_true_t-y_pred_t))
        cs_div = self.cs_div()       
        final_loss = mse + self.regularization_factor * cs_div
        return final_loss

Model:
class modelMSE_CSD (nn.Module):
    def __init__(self, input_size, bottleneck_dim):
        super(modelMSE_CSD, self).__init__()

        self.bottleneck_dim = bottleneck_dim
        self.input_size = input_size
        
        self.encoder = nn.Sequential(
            nn.Linear(self.input_size, 800),
            nn.ReLU(),
            nn.Linear(800,400),
            nn.ReLU(),
            nn.Linear(400, self.bottleneck_dim)
        )
        
        self.decoder = nn.Sequential(
            nn.Linear(self.bottleneck_dim, 400),
            nn.ReLU(),
            nn.Linear(400,800),
            nn.ReLU(),
            nn.Linear(800, self.input_size),
            nn.Sigmoid()
        )
        
    def forward(self,x):
        e_out = self.encoder(x)
        d_out = self.decoder(e_out)
        return e_out, d_out

Training Loop:
losses = []
for epoch in range(EPOCHS):
    running_loss = 0.0
    
    for batch_num, (data, target) in enumerate(dataloader):
        
        data = Variable(data, requires_grad = True)
        target = Variable(target, requires_grad = True)

        model_csd.train()
        latent_code, decoded_layer = model_csd(data)
        
        loss = MSE_ITL_RegLoss(latent_code=latent_code, prior=PRIOR_DIST, 
                               regularization_factor=REGULARIZER, kernel=KERNEL, batch_size=BATCH_SIZE)

        optimizer.zero_grad()
        **loss.backward()**
        optimizer.step()
        
        running_loss += loss.item() * data.size[0]
    losses.append(running_loss/len(dataset))
    print("Epoch: ",epoch,"Loss :", running_loss/len(dataset))```



Please suggest if the other method of declaring the custom loss is more easier.
Any links/resources would be appreciated.