checkpoints are not saved after implementing a custom loss 
Author: yhifny
Pull request: no

I subclassed Trainer class to implement the supervised constructive loss function (code found using google search). The code is running but it does not save the models checkpoints
`
class SupervisedContrastiveLoss(nn.Module):
    def __init__(self, temperature, device):
        """
        Implementation of the loss described in the paper Supervised Contrastive Learning :
        https://arxiv.org/abs/2004.11362
        :param temperature: int
        """
        super(SupervisedContrastiveLoss, self).__init__()
        self.temperature = temperature
        self.device = device

    def forward(self, projections, targets):
        """
        :param projections: torch.Tensor, shape [batch_size, projection_dim]
        :param targets: torch.Tensor, shape [batch_size]
        :return: torch.Tensor, scalar
        """
        projections = F.normalize(projections, p=2, dim=1)
        dot_product_tempered = torch.mm(projections, projections.T) / self.temperature
        # Minus max for numerical stability with exponential. Same done in cross entropy. Epsilon added to avoid log(0)
        exp_dot_tempered = (
            torch.exp(dot_product_tempered - torch.max(dot_product_tempered, dim=1, keepdim=True)[0]) + 1e-5
        )

        mask_similar_class = (targets.unsqueeze(1).repeat(1, targets.shape[0]) == targets).to(self.device)
        mask_anchor_out = (1 - torch.eye(exp_dot_tempered.shape[0])).to(self.device)
        mask_combined = mask_similar_class * mask_anchor_out
        cardinality_per_samples = torch.sum(mask_combined, dim=1)

        log_prob = -torch.log(exp_dot_tempered / (torch.sum(exp_dot_tempered * mask_anchor_out, dim=1, keepdim=True)))
        supervised_contrastive_loss_per_sample = torch.sum(log_prob * mask_combined, dim=1) / cardinality_per_samples
        supervised_contrastive_loss = torch.mean(supervised_contrastive_loss_per_sample)

        return supervised_contrastive_loss

class SCLTrainer(Trainer):

    def __init__(self,temperature, loss_weight, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.add_callback(EarlyStoppingCallback())
        self.temperature = temperature
        self.loss_weight = loss_weight
        self.device = self.args.device
        print ("SCL loss_weight: ", self.loss_weight)
        print ("SCL temperature: ", self.temperature)



    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        feature_vectors = outputs.hidden_states[-1] [:,0,:]
        logits = outputs["logits"] if isinstance(outputs, dict) else outputs[0]

        #print ("feature_vectors.shape: ", feature_vectors.shape)
        #print ("logits.shape: ", logits.shape)
        self.ce_loss  = nn.CrossEntropyLoss().to(self.device)
        self.scl_loss = SupervisedContrastiveLoss(self.temperature,self.device).to(self.device)

        loss = (1-self.loss_weight) * self.ce_loss(logits, labels) + self.loss_weight * self.scl_loss(feature_vectors, labels)

        return (loss, outputs) if return_outputs else loss
`
can you please let me know where is the problem?

Comments:
Hello! Could you provide the command you ran to run your script? If it's not an official example script, do you have the code you used handy? Thanks!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.