Defining loss_function

Loss Function

def calculate_loss(reconstructed, caption_prob, images, captions_transformed, mean, log_std):
    
    size = captions_transformed.shape[0]
    reconstruction_error = criterion(reconstructed, images)
    likelihoods = torch.stack([
        caption_prob[i, np.arange(MAX_CAPTION_LEN), captions_transformed[i]] for i in range(size)])
    
    
    log_likelihoods = -torch.log(likelihoods).sum()
    KL_divergence = - (1 - mean.pow(2) - torch.exp(2 * log_std) + (2 *log_std)).sum()
    
    return reconstruction_error + (log_likelihoods) + KL_divergence, log_likelihoods