RuntimeError: Placeholder storage has not been allocated on MPS device!

I am trying to train a CNN on a dataset loaded from local storage:

trainset = datasets.ImageFolder(root=train_dir, transform=train_transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)

Note: Problem may be caused by the 4 lines at the bottom here. It didn’t crash with VGG16 after I removed those lines, but I need the code for outputing one of 22 labels

model = models.resnet18(weights='IMAGENET1K_V1')
model.to(device=device)

for param in model.parameters():
  param.requires_grad = False
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs,22)

I have set the device to be mps:0 (I have also tried mps and confirmed my laptop has mps)

device = torch.device("mps:0" if torch.backends.mps.is_available() else "cpu")

In the training loop, I send the inputs to the device (I have also tried images.to(device) instead of images = …):

for epoch in range(1):
  train_loss = train(model, trainloader, optimizer, criterion)
def train(model,trainloader,optimizer,criterion):
  model.train()

  total_loss = []
  for data in trainloader:
    images, labels = data
    images = images.to(device)
    labels = labels.to(device)
    optimizer.zero_grad()

    outputs = model(images)
    ...

The code fails at the outputs = model(images) step

Error stack:

File /opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
...
File /opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: Placeholder storage has not been allocated on MPS device!

Ha… the issue is the missing device=device in nn.Linear()