# Using weights with transformers huggingface - running on GPUs

**URL:** <https://discuss.huggingface.co/t/using-weights-with-transformers-huggingface-running-on-gpus/16103>\
**Category:** 🤗Transformers\
**Created:** [March 26, 2022, 2:51am UTC](https://discuss.huggingface.co/t/using-weights-with-transformers-huggingface-running-on-gpus/16103 "2022-03-26T02:51:56Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![nitempe](https://avatars.discourse-cdn.com/v4/letter/n/54ee81/32.png) [@nitempe](https://discuss.huggingface.co/u/nitempe)\
**Post date:** [March 26, 2022, 2:51am UTC](https://discuss.huggingface.co/t/using-weights-with-transformers-huggingface-running-on-gpus/16103/1 "2022-03-26T02:51:56Z")

</div>

I came across this [tutorial](https://github.com/jlealtru/website_tutorials/blob/main/notebooks/Longformer%20with%20IMDB.ipynb) which performs Text classification with the Longformer. I came across this two links - [one](https://discuss.huggingface.co/t/how-can-i-use-class-weights-when-training/1067) and [two](https://huggingface.co/docs/transformers/master/main_classes/trainer) which talk about using class weights when the data is unbalanced.

```
# instantiate the trainer class and check for available devices
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=train_data,
    eval_dataset=test_data
)

```

The code from the first link is above

The code from the next 2 links is below

```
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 (suppose one has 3 labels with different weights)
        loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0, 3.0]))
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

```

I modified that code below as I am going to use GPU

```
device = 'cuda' if torch.cuda.is_available() else 'cpu'
#device = 'cpu'
device

```

Returns `'cuda'`

```
#class weights

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 (suppose one has 3 labels with different weights)
        loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 0.5243])).to(device)
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1)).to(device)
        return (loss, outputs) if return_outputs else loss

```

**my question)** I added part `.to(device)` to the below lines. Is that the correct way to modify the code so the entire code runs on GPUs without problems?

```
loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 0.5243])).to(device)
loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1)).to(device)

```