# Device error when extending Trainer class

**URL:** <https://discuss.huggingface.co/t/device-error-when-extending-trainer-class/36045>\
**Category:** 🤗Transformers\
**Created:** [April 9, 2023, 2:36pm UTC](https://discuss.huggingface.co/t/device-error-when-extending-trainer-class/36045 "2023-04-09T14:36:39Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![gianluca-db0](https://avatars.discourse-cdn.com/v4/letter/g/a88e4f/32.png) [@gianluca-db0](https://discuss.huggingface.co/u/gianluca-db0)\
**Post date:** [April 9, 2023, 2:36pm UTC](https://discuss.huggingface.co/t/device-error-when-extending-trainer-class/36045/1 "2023-04-09T14:36:39Z")

</div>

Hi everyone,

I have an unbalanced dataset, so I’ve been trying to extend the Trainer class as specified on Trainer docs\* and in [this thread](https://discuss.huggingface.co/t/how-can-i-use-class-weights-when-training/1067).  
So I have the following code:

```auto
import torch
import transformers
from torch import nn
from transformers import Trainer

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        # forward pass
        outputs = model(**inputs)
        logits = outputs.get('logits')
        # compute custom loss
        loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([0.2, 0.3, 0.5]))
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

```

I then define the trainer and try to train it in the following way:

```auto
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
tokenized_train = tokenized_train.with_format("torch")
tokenized_val = tokenized_val.with_format("torch")
model.to('cuda')

training_args = TrainingArguments(
    output_dir="my_awesome_model",
    learning_rate=1e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=5,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    #push_to_hub=True,
)

trainer = CustomTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

```

As you can see I also set the dataset format to ‘torch’ as specified in another thread\*, but I still get the following error at `trainer.train()`:

```auto
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument weight in method wrapper_CUDA_nll_loss_forward)

```

I’ve already tried to update the transformer library as specified in [here](https://discuss.huggingface.co/t/sending-a-dataset-or-datasetdict-to-a-gpu/17208) but it didn’t work. When I run this code using just `Trainer` class it works correctly, it throws this error only when trying to use this simple custom class.  
What could be the cause of this error and how could I fix it?

\*I would have added links to these but as a new user I can add only two links, sorry for the inconvenience.

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](https://discuss.huggingface.co/u/sgugger)\
**Post date:** [April 10, 2023, 1:28pm UTC](https://discuss.huggingface.co/t/device-error-when-extending-trainer-class/36045/2 "2023-04-10T13:28:18Z")

</div>

> [@gianluca-db0](#):
>
> `loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([0.2, 0.3, 0.5]))`

in this line the weights should be put on the same device as the model.