conv1 = nn.Conv2d(3,8,kernel_size=(3,4), padding='same')
conv2 = nn.Conv2d(8, 32, kernel_size=2, padding='same')
conv3 = nn.Conv2d(32,64,kernel_size=3, padding='same')
conv4 = nn.Conv2d(64, 128, kernel_size=3, padding='same')
conv5 = nn.Conv2d(128,256, kernel_size=3, padding='same')
conv6 = nn.Conv2d(256,256, kernel_size=(3,2), padding='same')

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = conv1
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = conv2
        self.conv3 = conv3
        self.conv4 = conv4
        self.conv5 = conv5
        self.conv6 = conv6
        self.fc1 = nn.Linear(256, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):


        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = self.pool(F.relu(self.conv4(x)))
        x = self.pool(F.relu(self.conv5(x)))
        x = x.view(-1, 256)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = torch.sigmoid(self.fc3(x))
        return x
cifar_model = Net()

training_data = datasets.CIFAR10(root="data", train=True, download=True,
                                  transform=transforms.Compose([
                                      transforms.ToTensor(),
                                      transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
                                  ]))

validation_data = datasets.CIFAR10(root="data", train=False, download=True,
                                  transform=transforms.Compose([
                                      transforms.ToTensor(),
                                      transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
                                  ]))

training_loader = DataLoader(training_data, 
                             batch_size=16, 
                             shuffle=True,
                             pin_memory=True)

validation_loader = DataLoader(validation_data,
                               batch_size=16,
                               shuffle=True,
                               pin_memory=True)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(cifar_model.parameters(), lr=1, momentum=0.9)
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

model = train_model(cifar_model, criterion, optimizer, exp_lr_scheduler,
                       num_epochs=25)

The train_model is exactly like above