Deepspeed trainer and custom loss weights
Hi all, I wrote a custom loss as suggested in this forum:loss_fct = torch.nn.CrossEntropyLoss(weight=class_weights_pt)
class SentTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        # compute custom loss
        labels = inputs.get("labels")
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else lossHowever, when integrating deepspeed, it complained about weight tensor not on the same device. How should I fix this issue?
2024-02-28T16:32:56Z
Morphe:loss_fct = torch.nn.CrossEntropyLoss(weight=class_weights_pt)I think you can defineloss_fctinside__init__function and thenset class_weights_ptto the appropriate device.def __init__(self, *args, class_weights: Optional[FloatTensor] = None, **kwargs):
        super().__init__(*args, **kwargs)
        if class_weights is not None:
            class_weights = class_weights.to(self.accelerator.device)
        self.loss_fct = torch.nn.CrossEntropyLoss(weight=class_weights)