Pytorch custom model automatically stored in cuda
I built a custom NN model like so:
class MyNNet(torch.nn.Module):
  
  def __init__(self, inp_dim, n_classes):
    super(MyNNet, self).__init__()
    self.flat = torch.nn.Flatten()
    self.l1 = torch.nn.Linear(inp_dim * inp_dim, 32)
    self.l2 = torch.nn.Linear(32, 16)
    self.l3 = torch.nn.Linear(16, n_classes)
  
  def forward(self, X):
    out = self.flat(X)
    out = F.relu(self.l1(out))
    out = F.relu(self.l2(out))
    return self.l3(out)

And a simple training script that updates the model parameters:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MyNNet(28, 10)
model.to(device)

optimizer = torch.optim.Adam(model.parameters())
loss = torch.nn.CrossEntropyLoss()

epochs = 20

for e in range(epochs):

  train_l = 0.
  for i, (s, c) in enumerate(train_loader):
    
    s.to(device)
    c.to(device)
    y_hat = model(s)

    l = loss(y_hat, c)
    train_l += l

    l.backward()
    optimizer.step()
    optimizer.zero_grad()

  print(f'Epoch: {e}, AvgLoss: {train_l / len(train_loader)}')

As in the script I store the model to cuda and so I do with each batch of the dataset (MNIST). However the folllowing error appears: Expected all tensors to be on the same device, but found at least two devices
but when I comment model.to(device), then the script works.  Does this mean PyTorch stores the custom models automatically into cuda?
Thanks.

Unlike Modules (where .to(...) works in-place), when moving Tensors to a device, you need to reassign them:
s = s.to(device)
c = c.to(device)

