Sorry, @ptrblck that was a blunder I made. I corrected and still the same error. I am new to PyTorch and DL please bear with me.
import torch.optim as optim
optimizer = optim.SGD(model.parameters(), lr = 1e-3)
criterion = nn.CrossEntropyLoss()
model = model.to(device)
criterion = model.to(device)