I’m trying to train a model for image classification. However, the training just doesn’t take place.
class CNN(nn.Module):
  def __init__(self):
    super(CNN, self).__init__()
    self.conv1 = nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2)
    self.conv2 = nn.Conv2d(64, 192, kernel_size=7, padding=2)
    self.conv3 = nn.Conv2d(192, 256, kernel_size=5, padding=1)

    self.fc1 = nn.Linear(256*9*9, 4096)
    self.fc2 = nn.Linear(4096, 1024)
    self.fc3 = nn.Linear(1024, 102)

  def forward(self,x):
    x = self.conv1(x)
    x = F.max_pool2d(x, kernel_size=3, stride=2)
    x = self.conv2(x)
    x = F.max_pool2d(x, kernel_size=3, stride=2)
    x = self.conv3(x)
    x = F.max_pool2d(x, kernel_size=3, stride=2)
    
    x = x.view(x.shape[0], -1)

    x = self.fc1(x)
    x = F.relu(x, inplace=True)
    x = self.fc2(x)
    x = F.relu(x, inplace=True)
    x = self.fc3(x)
    x = F.log_softmax(x, dim=1)

    return x

model = CNN()

criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters())
model.to(device)

NUM_EPOCHS = 10

for epoch in range(NUM_EPOCHS):
  train_loss = 0
  valloss = 0
  accuracy = 0

  counter = 0
  for inputs, labels in trainloader:
    # Move to device
    outputs= labels.to(device)
    inputs = inputs.to(device)

    # Get loss
    optimizer.zero_grad()
    preds = torch.exp(model(inputs))
    loss = criterion(preds, outputs)

    # Backprop
    loss.backward()
    optimizer.step()

    train_loss += loss.item()*inputs.size(0)
  
  # Validation
  with torch.no_grad():
    for inputs, labels in validationloader:
            # Move to device
            inputs, labels = inputs.to(device), labels.to(device)
            output = model(inputs)
            valloss = criterion(output, labels)
            valloss += valloss.item()*inputs.size(0)
            output = torch.exp(output)
            
            # Calculate accuracy
            top_p, top_class = output.topk(1, dim=1)
            equals = top_class == labels.view(*top_class.shape)
            accuracy += torch.mean(equals.type(torch.FloatTensor)).item()
    
    # Calculate loss
    train_loss = train_loss/len(trainloader.dataset)
    valloss = valloss/len(validationloader.dataset)
    
    # Print info
    print('Accuracy: ', accuracy/len(validationloader))
    print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(epoch, train_loss, valloss))