Repository: pytorch
Issue: "Expected a 'mps:0' generator device but found 'cpu'" using shuffle=True on DataLoader
Labels: module: dataloader,triaged,module: mps

### 🐛 Describe the bug

Using `shuffle=True` when creating a `DataLoader` results in some errors with generator type on macOS with MPS.

```python
import torch
import torch.utils.data as data

mps_device = torch.device("mps:0")
torch.set_default_device(mps_device)

x = torch.arange(100)
print(x)


class XDataset(data.Dataset):
    def __init__(self, x: torch.Tensor):
        super().__init__()
        self.x = x

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx: int):
        return self.x[idx]


dataset = XDataset(x)

data_loader = data.DataLoader(dataset)
print(next(iter(data_loader)))

data_loader = data.DataLoader(dataset, batch_size=8)
print(next(iter(data_loader)))

data_loader = data.DataLoader(dataset, batch_size=8, shuffle=True)
print(next(iter(data_loader)))
```

I get the following output:

```
tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
        18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
        36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
        54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
        72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89,
        90, 91, 92, 93, 94, 95, 96, 97, 98, 99], device='mps:0')
tensor([0], device='mps:0')
tensor([0, 1, 2, 3, 4, 5, 6, 7], device='mps:0')
```

Then the error:

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File ~/Code/toto-pytorch/mps_dataloader_shuffle_error.py:34
     31 print(next(iter(data_loader)))
     33 data_loader = data.DataLoader(dataset, batch_size=8, shuffle=True)
---> 34 print(next(iter(data_loader)))
     36 data_loader = data.DataLoader(
     37     dataset, batch_size=8, shuffle=True, generator=mps_device)
     38 print(next(iter(data_loader)))

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:630, in _BaseDataLoaderIter.__next__(self)
    627 if self._sampler_iter is None:
    628     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    629     self._reset()  # type: ignore[call-arg]
--> 630 data = self._next_data()
    631 self._num_yielded += 1
    632 if self._dataset_kind == _DatasetKind.Iterable and \
    633         self._IterableDataset_len_called is not None and \
    634         self._num_yielded > self._IterableDataset_len_called:

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:673, in _SingleProcessDataLoaderIter._next_data(self)
    672 def _next_data(self):
--> 673     index = self._next_index()  # may raise StopIteration
    674     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    675     if self._pin_memory:

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:620, in _BaseDataLoaderIter._next_index(self)
    619 def _next_index(self):
--> 620     return next(self._sampler_iter)

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/sampler.py:283, in BatchSampler.__iter__(self)
    281 batch = [0] * self.batch_size
    282 idx_in_batch = 0
--> 283 for idx in self.sampler:
    284     batch[idx_in_batch] = idx
    285     idx_in_batch += 1

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/sampler.py:165, in RandomSampler.__iter__(self)
    163 else:
    164     for _ in range(self.num_samples // n):
--> 165         yield from map(int, torch.randperm(n, generator=generator).numpy())
    166     yield from map(int, torch.randperm(n, generator=generator)[:self.num_samples % n].numpy())

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/_device.py:77, in DeviceContext.__torch_function__(self, func, types, args, kwargs)
     75 if func in _device_constructors() and kwargs.get('device') is None:
     76     kwargs['device'] = self.device
---> 77 return func(*args, **kwargs)

RuntimeError: Expected a 'mps:0' generator device but found 'cpu'
```

If I try to specify an MPS generator for the `DataLoader`:

```python
data_loader = data.DataLoader(dataset, batch_size=8, shuffle=True, generator=torch.Generator(device=mps_device))
print(next(iter(data_loader)))
```

I get a similar but slightly different error:

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[3], line 1
----> 1 print(next(iter(data_loader)))

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:630, in _BaseDataLoaderIter.__next__(self)
    627 if self._sampler_iter is None:
    628     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    629     self._reset()  # type: ignore[call-arg]
--> 630 data = self._next_data()
    631 self._num_yielded += 1
    632 if self._dataset_kind == _DatasetKind.Iterable and \
    633         self._IterableDataset_len_called is not None and \
    634         self._num_yielded > self._IterableDataset_len_called:

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:673, in _SingleProcessDataLoaderIter._next_data(self)
    672 def _next_data(self):
--> 673     index = self._next_index()  # may raise StopIteration
    674     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    675     if self._pin_memory:

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:620, in _BaseDataLoaderIter._next_index(self)
    619 def _next_index(self):
--> 620     return next(self._sampler_iter)

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/sampler.py:283, in BatchSampler.__iter__(self)
    281 batch = [0] * self.batch_size
    282 idx_in_batch = 0
--> 283 for idx in self.sampler:
    284     batch[idx_in_batch] = idx
    285     idx_in_batch += 1

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/data/sampler.py:165, in RandomSampler.__iter__(self)
    163 else:
    164     for _ in range(self.num_samples // n):
--> 165         yield from map(int, torch.randperm(n, generator=generator).numpy())
    166     yield from map(int, torch.randperm(n, generator=generator)[:self.num_samples % n].numpy())

File /opt/homebrew/lib/python3.11/site-packages/torch/utils/_device.py:77, in DeviceContext.__torch_function__(self, func, types, args, kwargs)
     75 if func in _device_constructors() and kwargs.get('device') is None:
     76     kwargs['device'] = self.device
---> 77 return func(*args, **kwargs)

RuntimeError: Expected a 'mps:0' generator device but found 'mps'
```

### Versions

Collecting environment information...
PyTorch version: 2.2.0.dev20231008
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.5.2 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.5 (main, Aug 24 2023, 15:09:45) [Clang 14.0.3 (clang-1403.0.22.14.1)] (64-bit runtime)
Python platform: macOS-13.5.2-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.1
[pip3] torch==2.2.0.dev20231008
[pip3] torchaudio==2.2.0.dev20231008
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.17.0.dev20231008
[conda] Could not collect

cc @SsnL @VitalyFedyunin @ejguan @dzhulgakov @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev