# Define model on cpu or gpu
model = CNN().to(device)

# Loss and optimizer
loss = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)