Hi, This is my code for the CIFAR-10 dataset:
Just sharing the code from the trainloader step. If you are still stuck I can share the entire notebook.
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)
valloader = torch.utils.data.DataLoader(valset, batch_size=64, shuffle=False, num_workers=2)

dataiterator = iter(trainloader)
samples = next(dataiterator)

# making the architecture
class Architecture(nn.Module):
    
    def __init__(self, input_num_channels, num_classes):
        super(Architecture, self).__init__()
        
        self.conv_blocks = nn.Sequential(

            nn.Conv2d(in_channels=input_num_channels, out_channels=32, kernel_size=(3,3), padding=(1,1), stride=(2,2), dilation=(2,2), bias=True),
            nn.BatchNorm2d(num_features=32),
            nn.LayerNorm([32,15,15], elementwise_affine=True),
            nn.InstanceNorm2d(num_features=32, affine=True),
            nn.GroupNorm(num_groups=4, num_channels=32, affine=True),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2,2), padding=(0,0), stride=(1,1), dilation=(2,2))
            
        )
        
        self.linear_blocks = nn.Sequential(

            nn.Linear(in_features=13*13*32, out_features=2000, bias=True),
            nn.BatchNorm1d(num_features=2000),
            nn.LayerNorm([2000], elementwise_affine=True),
            nn.GroupNorm(num_groups=100, num_channels=2000, affine=True),
            nn.Tanh(),
            nn.Linear(in_features=2000, out_features=num_classes, bias=True)
        )
        
    def forward(self, x):
        
        x = self.conv_blocks(x)
        x = x.view(x.size(0), -1)
        x = self.linear_blocks(x)
        
        return x
    


def weight_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.xavier_normal_(m.weight)
#         nn.init.kaiming_normal_(m.bias)
        
    elif isinstance(m, (nn.Linear, nn.BatchNorm1d)):
        nn.init.normal_(m.weight)
        nn.init.normal_(m.bias)
        
        
net = Architecture(3, 10)
net.apply(weight_init)


criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(tenet.parameters(), lr=0.001, momentum=0.9)

num_epochs=10

for epoch in range(num_epochs):
    
    print("Epoch:", epoch+1)
    
    running_loss = 0.0
    
    
    print(optimizer.param_groups[0]['lr'])
    
    
    for index, data in enumerate(trainloader):

        inputs, labels = data
        
        optimizer.zero_grad()
        
        outputs = tenet(inputs)
#         print(outputs.shape)
#         print(labels.shape)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        print(loss.item())
        
        running_loss += loss.item()
    
    scheduler.step()
    
    print("training loss after epoch", epoch+1, ":", running_loss/len(trainloader))
    
    train_accuracy = calculate_accuracy(trainloader)
    print("train accuracy after epoch", epoch+1, ":", train_accuracy)
    
    val_accuracy = calculate_accuracy(valloader)
    print("val accuracy after epoch", epoch+1, ":", val_accuracy)


This was probably the first code for me in pytorch. After that, I learned a lot.
That it is best to define a function for training and testing and then calling them in the for epoch loop.
Like:
def train(model: nn.Module,
          iterator: torch.utils.data.DataLoader,
          optimizer: optim.Optimizer,
          criterion: nn.Module, 
          clip_value):

    model.train()

    epoch_loss = 0

    for _, (imgs, texts, text_seq_lens_list, targets) in enumerate(iterator):

        optimizer.zero_grad()

        output = model(imgs, texts, text_seq_lens_list)

        loss = criterion(output, targets)

        loss.backward()

        nn.utils.clip_grad_norm_(model.parameters(), clip_value)

        optimizer.step()

        epoch_loss += loss.item()
        print(loss.item())

    return epoch_loss / len(iterator)

num_epochs = 10
clip_value = 1.0

for epoch in range(num_epochs):

    print(epoch)
    train_loss = train(model, trainloader, optimizer, criterion, clip_value)
    print("train loss:", train_loss)


I didn’t write this piece, I found it somewhere on github. I think on Ben Trevett’s github repo.