Tensors' device passed to a model is not correct when ACCELERATE_TORCH_DEVICE is privateuseone

[kiszk]

System Info

• transformers version: 4.43.0
• Platform: Linux-5.4.0-187-generic-x86_64-with-glibc2.31
• Python version: 3.10.14
• Huggingface_hub version: 0.23.4
• Safetensors version: 0.4.3
• Accelerate version: 0.31.0
• Accelerate config: not found
• PyTorch version (GPU?): 2.3.1+cu121 (True)
• Tensorflow version (GPU?): 2.15.1 (True)
• Flax version (CPU?/GPU?/TPU?): 0.7.0 (cpu)
• Jax version: 0.4.13
• JaxLib version: 0.4.13
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: Tesla V100-PCIE-32GB

Who can help?

@muellerzr @SunMarc

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

Execute the following script test.py with accelerators(TPU, CUDA, CPU, privateuseone:0) with accelerate.

import sys
import accelerate
import torch
import transformers
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
)

model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def forward(input_ids=None, attention_mask=None, **kwargs):
    # print devices for input tensors instead of doing inference
    print(input_ids.device, attention_mask.device)
    sys.exit(1)
model.forward = forward

trainer = Trainer(
    model = model,
    tokenizer = tokenizer,
)

class MyDataset(torch.utils.data.Dataset):
    def __init__(self):
        self.data_list = [{"input_ids":[1], "label":[1]}]

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, idx):
        return self.data_list[idx]

predict_dataset = MyDataset()
trainer.predict(predict_dataset)

Expected behavior

When ACCELERATE_TORCH_DEVICE=privateuseone:0 python test.py is executed, I expect privateuseone:0 privateuseone:0 will be displayed. However, I saw cpu cpu.

With TPU, I got xla:0 xla:0. With NVIDIA GPU, I got cuda:0 cuda:0. With only CPU, I got cpu cpu.