feedback loop: Expected a 'cpu' device type for generator but found 'mps'

Question

I am trying to finetune a HuggingFace model on MacOS M2. However, the torch operator 'aten::random_' does not exist yet for MPS. Therefore, I ran the program with

PYTORCH_ENABLE_MPS_FALLBACK=1 

Then I got the following message "Expected a 'cpu' generator device type but found 'mps'". I tried to fix that by specifiying generator=torch.Generator(device='mps') in the dataloader. This however turns the message around and yields "Expected a 'mps:0' generator device but found 'mps'". It therefore seems I'm stuck in a loop.

This is the full code I am using

from datasets import load_dataset
from transformers import AutoTokenizer
import torch
from torch.utils.data import DataLoader

dataset = load_dataset('yelp_review_full')
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

tokenized_datasets = tokenized_datasets.remove_columns(["text"])
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format("torch")
small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))


train_dataloader = DataLoader(small_train_dataset, shuffle=True, batch_size=8, generator = torch.Generator(device='cpu').manual_seed(int(1)))
eval_dataloader = DataLoader(small_eval_dataset, batch_size=8, generator=torch.Generator(device='cpu'))
for batch in train_dataloader:
    print(batch)
    # of course I want to do other stuff here but opening the loop already produces the error. 

I am using torch 2.0.1, datasets 2.14.0 together with transformers 4.31.0