Solved by ptrblck in post #6 
              
              
                Your general workflow should work as shown in this small example: 
# Setup
model = models.resnet18()
data = torch.randn(1, 3, 224, 224)
target = torch.randint(0, 1000, (1,))
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# Train
optimizer.zero_grad()
out …