I am beginner in PyTorch; with the help of a YouTube channel “Deep Lizard” i learned about the tensors along with the operations. Now, i am working on a vehicle classification model, having 10 different classes of 1000 images each. But my model is not converging the training loss. I have attached the some chunks of code below. Please let me know, where i did mistake. Thanks.
#Resize Image
transform_in = transforms.Compose([
transforms.Resize([128, 128]),
transforms.ToTensor(),
transforms.Normalize((0.4915, 0.4823, 0.4468),
(0.2470, 0.2435, 0.2616))
])
#Model Class
class Net(nn.Module):
def init(self):
super().init()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
self.pool1 = nn.MaxPool2d(2,2)
    self.conv2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
    self.pool2 = nn.MaxPool2d(2,2)
    
    self.conv3 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
    self.pool3 = nn.MaxPool2d(2,2)
    
    self.conv4 = nn.Conv2d(16, 8, kernel_size=3, padding=1)
    self.pool4 = nn.MaxPool2d(2,2)
    
    self.fc1 = nn.Linear(8 * 8 * 8, 288)
    self.fc2 = nn.Linear(288, 144)
    self.fc3 = nn.Linear(144, 72)
    self.fc4 = nn.Linear(72, 2)
    
def forward(self, x):
    out = F.max_pool2d(torch.relu(self.conv1(x)), 2)
    out = F.max_pool2d(torch.relu(self.conv2(out)), 2)
    out = F.max_pool2d(torch.relu(self.conv3(out)), 2)
    out = F.max_pool2d(torch.relu(self.conv4(out)), 2)
    out = out.view(-1, 8 * 8 * 8)
    out = torch.relu(self.fc1(out))
    out = torch.relu(self.fc2(out))
    out = torch.relu(self.fc3(out))
    out = torch.relu(self.fc4(out))        
    out = F.softmax(out, dim=1)
    return out

#Training
def training_loop(n_epochs, optimizer, model, loss_fn, train_loader):
for epoch in range(1, n_epochs + 1):
loss_train = 0.0
for imgs, labels in train_loader:
imgs = imgs.to(device=device)  # <1>
labels = labels.to(device=device)
outputs = model(imgs)
loss = loss_fn(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        loss_train += loss.item()

    if epoch == 1 or epoch % 2 == 0:
        print('Epoch {}, Training loss {}' epoch, loss_train / len(train_loader)))

#Test
model = Net().to(device=device)  # <1>
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
loss_fn = nn.CrossEntropyLoss()
training_loop(
n_epochs = num_epochs,
optimizer = optimizer,
model = model,
loss_fn = loss_fn,
train_loader = train_loader,
)