Expected a 'mps:0' generator device but found 'mps'

Fyu: <p>I’m trying to load custom data for a CNN via mps on a MacBook pro M3 pro but encounter the issue where the generator expects a mps:0 generator but gets mps</p>
<ul>
<li>Python ver: 3.11.6</li>
<li>PyTorch ver: 2.1.1</li>
<li>PyVision ver: 0.16.1</li>
<li>Environment: Jupyter Notebook (on VSCode)</li>
</ul>
<p>Code:</p>
<p>if torch.backends.mps.is_available():<br>
mps_device = torch.device(‘mps:0’)<br>
torch.set_default_device(mps_device)<br>
else:<br>
print (‘MPS device not found.’)</p>
<p>trainset = datasets.ImageFolder(root=train_dir, transform=train_transform)<br>
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, generator=torch.Generator(device=mps_device))</p>
<p>model = models.resnet18(weights=‘IMAGENET1K_V1’)<br>
…</p>
<p>Error stack:</p>
<p>File <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:630" rel="noopener nofollow ugc">/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:630</a>, in _BaseDataLoaderIter.<strong>next</strong>(self) <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:627" rel="noopener nofollow ugc">627</a> if self._sampler_iter is None: <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:628" rel="noopener nofollow ugc">628</a> # TODO(<a href="https://github.com/pytorch/pytorch/issues/76750" class="inline-onebox" rel="noopener nofollow ugc">Bug in dataloader iterator found by mypy · Issue #76750 · pytorch/pytorch · GitHub</a>) <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:629" rel="noopener nofollow ugc">629</a> self._reset() # type: ignore[call-arg] → <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:630" rel="noopener nofollow ugc">630</a> data = self._next_data() <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:631" rel="noopener nofollow ugc">631</a> self._num_yielded += 1 <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:632" rel="noopener nofollow ugc">632</a> if self._dataset_kind == _DatasetKind.Iterable and \ <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:633" rel="noopener nofollow ugc">633</a> self._IterableDataset_len_called is not None and \ <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/data/dataloader.py:634" rel="noopener nofollow ugc">634</a> self._num_yielded &gt; self._IterableDataset_len_called:<br>
…<br>
<a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/_device.py:75" rel="noopener nofollow ugc">75</a> if func in _device_constructors() and kwargs.get(‘device’) is None: <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/_device.py:76" rel="noopener nofollow ugc">76</a> kwargs[‘device’] = self.device —&gt; <a href="https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/torch/utils/_device.py:77" rel="noopener nofollow ugc">77</a> return func(*args, **kwargs)<br>
RuntimeError: Expected a ‘mps:0’ generator device but found ‘mps’</p>

ptrblck: <p>Try to use the <code>mps</code> backend explicitly instead of using <code>set_default_device</code>.</p>