CNN on Fashion-MNIST dataset with Pytorch

Fashion_MNST.py

import argparse
import torch 
import torch.nn as nn 
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.optim
import torch.utils.data

CUDA = False

#----------------------------- NET CLASS -------------------------------#

class Net(nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    #defining the layers
    self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=5) # 28 X 28 --> 24 X 24 --> 12 X 12
    self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=5) #12 X 12 --> 8 X 8 --> 4 X 4  
    self.fc1 = nn.Linear(in_features=16 * 4 * 4, out_features=64) # 4 X 4 X 32 --> 64  
    self.fc2 = nn.Linear(in_features=64, out_features=10) # 64 --> 10

  #forward pass
  def forward(self, x):
    x = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=2, stride=2)
    x = F.max_pool2d(F.relu(self.conv2(x)), kernel_size=2, stride=2)
    x = F.max_pool2d(f.relu(self.conv3(x)), kernel_size=2, stride=2)
    x = x.view(-1, 16 * 4 * 4 )
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return x
  
def main(params):
  #get the train data download it if needed
  train = torchvision.datasets.FashionMNIST(
          root = './data/FashionMNIST',
          train = True,
          download = True,
          transform = transforms.Compose([ transforms.ToTensor()]))
  
  #get the test data download it if needed
  test = torchvision.datasets.FashionMNIST(
          root = './data/FashionMNIST',
          train = False,
          download = True,
          transform = transforms.Compose([ transforms.ToTensor()]))
  #devide the train data into batchs and shaffle them 
  trainset = torch.utils.data.DataLoader(trainset, batch_size = 16,
                                           shuffle = True, num_workers = 2)
  
  #devide the train data into batchs and shaffle them 
  testset = torch.utils.data.DataLoader(trainset, batch_size = 16,
                                           shuffle = True, num_workers = 2)
  
  net = Net() #creating the net 
  optimizer = torch.optim.SGD(net.parameters(), params.lr) #creating the optimizer 
  
  if CUDA: #IF CUDA is enabled transforme the model into the GPU 
        model = model.cuda()
  
  for i in range(params.epochs):
    print("=================\n=== EPOCH "+str(i+1)+" =====\n=================\n")
    for data in trainset:
      X, y = data
      net.zero_grad()
      output = net(X)
      loss = F.cross_entropy(preds, labels)
      loss.backward()
      optimizer.step()
    print(loss)