Repository: pytorch
Issue: RuntimeError using nested tensor in Apple M1 device MPS
Labels: triaged,module: nestedtensor,module: mps

### 🐛 Describe the bug

```
Traceback (most recent call last):
  File "/Users/aisiqg/Software/venv/pydptb/bin/dptb", line 8, in <module>
    sys.exit(main())
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/dptb/entrypoints/main.py", line 408, in main
    train(**dict_args)
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/dptb/entrypoints/train.py", line 223, in train
    trainer.run(trainer.train_options["num_epoch"])
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/dptb/nnops/base_trainer.py", line 51, in run
    self.epoch()
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/dptb/nnops/trainer.py", line 192, in epoch
    self.iteration(ibatch, next(iter(self.reference_loader)))
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/dptb/nnops/trainer.py", line 84, in iteration
    batch = batch.to(self.device)
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/dptb/utils/torch_geometric/data.py", line 322, in to
    self.apply(lambda x: x.to(device, **kwargs), *keys)
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/dptb/utils/torch_geometric/data.py", line 306, in apply
    self[key] = torch.nested.as_nested_tensor(item)
  File "/Users/aisiqg/Software/venv/pydptb/lib/python3.9/site-packages/torch/nested/__init__.py", line 58, in as_nested_tensor
    return torch._nested_tensor_from_tensor_list(tensor_list, dtype, None, device, None)
RuntimeError: storage_device.is_cpu() || storage_device.is_cuda() || storage_device.is_privateuseone() INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/NestedTensorImpl.cpp":185, please report a bug to PyTorch. NestedTensorImpl storage must be either CUDA, CPU or privateuseone but got mps:0

```

### Versions

Collecting environment information...
PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.4.1 (x86_64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.22.2
Libc version: N/A

Python version: 3.9.7 (default, Sep 16 2021, 08:50:36)  [Clang 10.0.0 ] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-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] flake8==6.1.0
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.24.4
[pip3] numpydoc==1.1.0
[pip3] torch==2.1.0
[pip3] torchaudio==2.1.0
[pip3] torchsort==0.1.7
[pip3] torchvision==0.16.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.4.0           hecd8cb5_637  
[conda] mkl-service               2.4.0            py39h9ed2024_0  
[conda] mkl_fft                   1.3.1            py39h4ab4a9b_0  
[conda] mkl_random                1.2.2            py39hb2f4e1b_0  
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] torch                     2.1.0                    pypi_0    pypi
[conda] torchaudio                2.1.0                    pypi_0    pypi
[conda] torchsort                 0.1.7                    pypi_0    pypi
[conda] torchvision               0.16.0                   pypi_0    pypi

cc @cpuhrsch @jbschlosser @bhosmer @drisspg @soulitzer @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen