RuntimeError: CUDA error: device-side assert triggered
### System Info

- `transformers` version: 4.11.3
- Platform: Linux-5.4.0-58-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 2.0.0+cu117 (True)
- Tensorflow version (GPU?): 2.10.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker @younesbelkada @sgugger 

I am trying out this [tutorial](https://flower.dev/docs/quickstart-huggingface.html) with 2 clients. In the first client, I am using the same dataset as that of the tutorial (IMDB) but for the 2nd client I am using [this](https://huggingface.co/datasets/sst2) dataset. While running the `server.py`, `client_1.py`, and `client_2.py` I am facing the following error while running `client_2.py`. I have also attached all the files below.

Error:
```
Traceback (most recent call last):
  File "client_2.py", line 140, in <module>
    main()
  File "client_2.py", line 136, in main
    fl.client.start_numpy_client(server_address="localhost:5040", client=IMDBClient())
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/app.py", line 208, in start_numpy_client
    start_client(
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/app.py", line 142, in start_client
    client_message, sleep_duration, keep_going = handle(
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/grpc_client/message_handler.py", line 70, in handle
    return _evaluate(client, server_msg.evaluate_ins), 0, True
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/grpc_client/message_handler.py", line 182, in _evaluate
    evaluate_res = client.evaluate(evaluate_ins)
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/app.py", line 276, in _evaluate
    results = self.numpy_client.evaluate(parameters, ins.config)  # type: ignore
  File "client_2.py", line 131, in evaluate
    loss, accuracy = test(net, testloader)
  File "client_2.py", line 97, in test
    loss += outputs.loss.item()
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

`client_1.py` file:
```
from collections import OrderedDict
import warnings

import flwr as fl
import torch
import numpy as np

import random
from torch.utils.data import DataLoader

from datasets import load_dataset, load_metric

from transformers import AutoTokenizer, DataCollatorWithPadding
from transformers import AutoModelForSequenceClassification
from transformers import AdamW

# import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

warnings.filterwarnings("ignore", category=UserWarning)
# DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

DEVICE = "cuda:0"

CHECKPOINT = "distilbert-base-uncased"  # transformer model checkpoint


def load_data():
    """Load IMDB data (training and eval)"""
    raw_datasets = load_dataset("imdb")
    raw_datasets = raw_datasets.shuffle(seed=42)

    # remove unnecessary data split
    del raw_datasets["unsupervised"]

    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)

    def tokenize_function(examples):
        return tokenizer(examples["text"], truncation=True)

    # random 100 samples
    # population = random.sample(range(len(raw_datasets["train"])), 100)

    tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
    # tokenized_datasets["train"] = tokenized_datasets["train"].select(population)
    # tokenized_datasets["test"] = tokenized_datasets["test"].select(population)

    tokenized_datasets = tokenized_datasets.remove_columns("text")
    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
    trainloader = DataLoader(
        tokenized_datasets["train"],
        shuffle=True,
        batch_size=32,
        collate_fn=data_collator,
    )

    testloader = DataLoader(
        tokenized_datasets["test"], batch_size=32, collate_fn=data_collator
    )

    return trainloader, testloader


def train(net, trainloader, epochs):
    optimizer = AdamW(net.parameters(), lr=5e-5)
    net.train()
    for i in range(epochs):
        print("Epoch: ", i+1)
        j = 1
        for batch in trainloader:
            print("####################### The batch number is: ", j)
            batch = {k: v.to(DEVICE) for k, v in batch.items()}
            outputs = net(**batch)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            j += 1


def test(net, testloader):
    metric = load_metric("accuracy")
    loss = 0
    net.eval()
    for batch in testloader:
        batch = {k: v.to(DEVICE) for k, v in batch.items()}
        with torch.no_grad():
            outputs = net(**batch)
        logits = outputs.logits
        loss += outputs.loss.item()
        predictions = torch.argmax(logits, dim=-1)
        metric.add_batch(predictions=predictions, references=batch["labels"])
    loss /= len(testloader.dataset)
    accuracy = metric.compute()["accuracy"]
    return loss, accuracy


def main():
    net = AutoModelForSequenceClassification.from_pretrained(
        CHECKPOINT, num_labels=2
    ).to(DEVICE)

    trainloader, testloader = load_data()

    # Flower client
    class IMDBClient(fl.client.NumPyClient):
        def get_parameters(self, config):
            return [val.cpu().numpy() for _, val in net.state_dict().items()]

        def set_parameters(self, parameters):
            params_dict = zip(net.state_dict().keys(), parameters)
            state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
            net.load_state_dict(state_dict, strict=True)

        def fit(self, parameters, config):
            self.set_parameters(parameters)
            print("Training Started...")
            train(net, trainloader, epochs=1)
            print("Training Finished.")
            return self.get_parameters(config={}), len(trainloader), {}

        def evaluate(self, parameters, config):
            self.set_parameters(parameters)
            loss, accuracy = test(net, testloader)
            print({"loss": float(loss), "accuracy": float(accuracy)})
            return float(loss), len(testloader), {"loss": float(loss), "accuracy": float(accuracy)}

    # Start client
    fl.client.start_numpy_client(server_address="localhost:5040", client=IMDBClient())


if __name__ == "__main__":
    main()
```

`client_2.py` file:
```
from collections import OrderedDict
import warnings

import flwr as fl
import torch
import numpy as np

import random
from torch.utils.data import DataLoader

from datasets import load_dataset, load_metric

from transformers import AutoTokenizer, DataCollatorWithPadding
from transformers import AutoModelForSequenceClassification
from transformers import AdamW
#from transformers import tokenized_datasets

# import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '2,3'

warnings.filterwarnings("ignore", category=UserWarning)
# DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

DEVICE = "cuda:1"

CHECKPOINT = "distilbert-base-uncased"  # transformer model checkpoint


def load_data():
    """Load IMDB data (training and eval)"""
    raw_datasets = load_dataset("sst2")
    # raw_datasets = load_dataset("yhavinga/imdb_dutch")
    raw_datasets = raw_datasets.shuffle(seed=42)

    # remove unnecessary data split
    del raw_datasets["validation"]
    # del raw_datasets["unsupervised"]

    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)

    def tokenize_function(examples):
        return tokenizer(examples["sentence"], truncation=True)

    # random 100 samples
    # population = random.sample(range(len(raw_datasets["train"])), 100)

    tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
    # tokenized_datasets["train"] = tokenized_datasets["train"].select(population)
    # tokenized_datasets["test"] = tokenized_datasets["test"].select(population)

    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
    
    tokenized_datasets = tokenized_datasets.remove_columns(["idx", "sentence"])
    
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
    trainloader = DataLoader(
        tokenized_datasets["train"],
        shuffle=True,
        batch_size=32,
        collate_fn=data_collator,
    )

    testloader = DataLoader(
        tokenized_datasets["test"], batch_size=32, collate_fn=data_collator
    )

    return trainloader, testloader


def train(net, trainloader, epochs):
    optimizer = AdamW(net.parameters(), lr=5e-4)
    net.train()
    for i in range(epochs):
        print("Epoch: ", i+1)
        j = 1
        # print("####################### The length of the trainloader is: ", len(trainloader))
        for batch in trainloader:
            print("####################### The batch number is: ", j)
            batch = {k: v.to(DEVICE) for k, v in batch.items()}
            outputs = net(**batch)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            j += 1


def test(net, testloader):
    metric = load_metric("accuracy")
    loss = 0
    net.eval()
    for batch in testloader:
        batch = {k: v.to(DEVICE) for k, v in batch.items()}
        with torch.no_grad():
            outputs = net(**batch)
        logits = outputs.logits
        loss += outputs.loss.item()
        predictions = torch.argmax(logits, dim=-1)
        metric.add_batch(predictions=predictions, references=batch["labels"])
    loss /= len(testloader.dataset)
    accuracy = metric.compute()["accuracy"]
    return loss, accuracy


def main():
    net = AutoModelForSequenceClassification.from_pretrained(
        CHECKPOINT, num_labels=2
    ).to(DEVICE)

    trainloader, testloader = load_data()

    # Flower client
    class IMDBClient(fl.client.NumPyClient):
        def get_parameters(self, config):
            return [val.cpu().numpy() for _, val in net.state_dict().items()]

        def set_parameters(self, parameters):
            params_dict = zip(net.state_dict().keys(), parameters)
            state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
            net.load_state_dict(state_dict, strict=True)

        def fit(self, parameters, config):
            self.set_parameters(parameters)
            print("Training Started...")
            train(net, trainloader, epochs=1)
            print("Training Finished.")
            return self.get_parameters(config={}), len(trainloader), {}

        def evaluate(self, parameters, config):
            self.set_parameters(parameters)
            loss, accuracy = test(net, testloader)
            print({"loss": float(loss), "accuracy": float(accuracy)})
            return float(loss), len(testloader), {"loss": float(loss), "accuracy": float(accuracy)}

    # Start client
    fl.client.start_numpy_client(server_address="localhost:5040", client=IMDBClient())


if __name__ == "__main__":
    main()
```

`server.py` file:
```
import flwr as fl
import torch

from collections import OrderedDict
from logging import WARNING
from flwr.common import (
    ndarrays_to_parameters,
    parameters_to_ndarrays,
)
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer
from flwr.server.strategy.aggregate import aggregate
from flwr.common.logger import log

from transformers import pipeline

CHECKPOINT = "distilbert-base-uncased"  # transformer model checkpoint
DEVICE = torch.device("cpu")
# DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class Strategy(fl.server.strategy.FedAvg):
    def aggregate_fit(
        self,
        server_round,
        results,
        failures,
    ):
        if not results:
            return None, {}
        # Do not aggregate if there are failures and failures are not accepted
        if not self.accept_failures and failures:
            return None, {}

        # Convert results
        weights_results = [
            (parameters_to_ndarrays(fit_res.parameters), fit_res.num_examples)
            for _, fit_res in results
        ]
        self.aggr_weights = aggregate(weights_results)
        parameters_aggregated = ndarrays_to_parameters(self.aggr_weights)

        # Aggregate custom metrics if aggregation fn was provided
        metrics_aggregated = {}
        if self.fit_metrics_aggregation_fn:
            fit_metrics = [(res.num_examples, res.metrics) for _, res in results]
            metrics_aggregated = self.fit_metrics_aggregation_fn(fit_metrics)
        elif server_round == 1:  # Only log this warning once
            log(WARNING, "No fit_metrics_aggregation_fn provided")
        return parameters_aggregated, metrics_aggregated

    
if __name__ == "__main__":
    
    def weighted_average(metrics):
        accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
        losses = [num_examples * m["loss"] for num_examples, m in metrics]
        examples = [num_examples for num_examples, _ in metrics]
        accuracy = sum(accuracies) / sum(examples)
        loss = sum(losses) / sum(examples)
        print("Accuracy: ", accuracy)
        print("Loss: ", loss)
        return {"accuracy": accuracy, "loss": loss}

    net = AutoModelForSequenceClassification.from_pretrained(
        CHECKPOINT, num_labels=2
    ).to(DEVICE)
    # Define strategy
    strategy = Strategy(
        min_fit_clients=1,
        min_evaluate_clients=1,
        min_available_clients=1,
        fraction_fit=1.0,
        fraction_evaluate=1.0,
        evaluate_metrics_aggregation_fn=weighted_average,
    )
    

    # Start server
    fl.server.start_server(
        server_address="localhost:5040",
        config=fl.server.ServerConfig(num_rounds=2),
        strategy=strategy,
    )
    params_dict = zip(net.state_dict().keys(), strategy.aggr_weights)
    state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
    net.load_state_dict(state_dict)
    classifier = pipeline("sentiment-analysis", model=net, tokenizer=AutoTokenizer.from_pretrained(CHECKPOINT))
    positive_1 = "That was amazing!!!"
    negative_1 = "I feel so sad about this..."
    positive_2 = "I liked it!!"
    negative_2 = "I hated it!!"
    print(positive_1, classifier(positive_1))
    print(negative_1, classifier(negative_1))
    print(positive_2, classifier(positive_2))
    print(negative_2, classifier(negative_2))

    # Dutch inference
    dutch_pos_1 = "ik vond de film leuk"
    dutch_neg_1 = "Ik haatte de film"
    print(dutch_pos_1, classifier(dutch_pos_1))
    print(dutch_neg_1, classifier(dutch_neg_1))
    
    torch.save(net.state_dict(), "/home/saurav/quickstart_huggingface/server_model.pt")
```

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Run the above provided files in 3 separate terminal windows.

### Expected behavior

The training should happen without causing an error.
This usually means there is bad indexing somewhere in your code. You should really bring up the issue with the persons who wrote the tutorials as it looks like there is a bug in their code. You can try to run the code on the CPU to see where the error stems from, or post us a minimal reproducer that doesn't use third-party libraries.<|||||>But `client_1.py` works well and there isn't much difference in both of their code, only the dataset is different.<|||||>I ran it on a CPU rather than a GPU and have got some more information about the error:

```
Traceback (most recent call last):
  File "client_2.py", line 140, in <module>
    main()
  File "client_2.py", line 136, in main
    fl.client.start_numpy_client(server_address="localhost:5040", client=IMDBClient())
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/app.py", line 208, in start_numpy_client
    start_client(
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/app.py", line 142, in start_client
    client_message, sleep_duration, keep_going = handle(
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/grpc_client/message_handler.py", line 70, in handle
    return _evaluate(client, server_msg.evaluate_ins), 0, True
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/grpc_client/message_handler.py", line 182, in _evaluate
    evaluate_res = client.evaluate(evaluate_ins)
  File "/home/saurav/.local/lib/python3.8/site-packages/flwr/client/app.py", line 276, in _evaluate
    results = self.numpy_client.evaluate(parameters, ins.config)  # type: ignore
  File "client_2.py", line 131, in evaluate
    loss, accuracy = test(net, testloader)
  File "client_2.py", line 95, in test
    outputs = net(**batch)
  File "/home/saurav/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/saurav/.local/lib/python3.8/site-packages/transformers/models/distilbert/modeling_distilbert.py", line 763, in forward
    loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
  File "/home/saurav/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/saurav/.local/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/saurav/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 3029, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
IndexError: Target -1 is out of bounds.
```<|||||>This probably means some of your labels are `-1` which is not a valid label. If you were attempting to put a fake label to pad a batch, -100 is the value for this in PyTorch.<|||||>I am using this dataset - https://huggingface.co/datasets/sst2 and I noticed that the test set has -1 value. How do I remove it?<|||||>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/main/CONTRIBUTING.md) are likely to be ignored.