Hi there. Kinda new to PyTorch. I have implemented a fully connected neural network and am trying to compute accuracy on the training samples. After the first epoch, I get a certain accuracy, after the second epoch, i get a slightly better accuracy, great. But then accuracy doesn’t change. Can anyone point me to the bug please? Note that I am not yet testing. I am still debugging my training algorithm
#HINT: note that your training time should not take many days.

max_epoch = 100
train_batch = 4
test_batch = 500
learning_rate = 0.1
use_gpu = torch.cuda.is_available()

def main(): # you are free to change parameters
    
    xl_file = 'labels.xlsx'
    root_dir = 'C:\\Users\\t.anjary\\Desktop\\ML\\ML HW3\\images'
    train_set, val_set, test_set = get_dataset(xl_file, root_dir)
    
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=train_batch, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_set, batch_size=test_batch, shuffle=False)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=test_batch, shuffle=False)
    
    model = FNet().float() #.cuda()
    criterion = torch.nn.MSELoss(reduction='sum')
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)#, momentum=0.9, weight_decay=5e-04) # you can play with momentum and weight_decay parameters as well
    
    for epoch in range(max_epoch):
        print('\n epoch '+ str(epoch))
        train(epoch, model, criterion, optimizer, train_loader)
    
''' Train your network for a one epoch '''
def train(epoch, model, criterion, optimizer, loader): # you are free to change parameters
    #print('train is called')
    accuracies =0
    model.train()
    correct = 0
    total = 0
    for batch_idx, (data, labels) in enumerate(loader):
        # TODO:
        # Implement training code for a one iteration
        
        #forward
        optimizer.zero_grad()
        data = data.float() #.cuda()
        labels = labels.float() #.cuda()
        y_pred = model(data)
        
        #loss
        loss = criterion(y_pred, labels)
        
        #update
        loss.backward()
        optimizer.step()
        
        #performance
        #y_pred = y_pred.cpu().numpy()
        y_pred = y_pred.flatten()
        total = labels.size(0)
        correct = (y_pred == labels).sum().item()
        accuracy = 100 * correct / total
        accuracies += accuracy 
        
        #info = 'iteration: ' + str(batch_idx) + '  loss: ' + str(loss.item()) + '  accuracy: ' + str(accuracy)
        #print(info)
        #sys.stdout.flush()
        sys.stdout.write('\r iteration: ' + str(batch_idx) + '  loss: ' + str(loss.item()) + '  accuracy: ' + str(accuracy))
        
        
    accuracies = accuracies/(batch_idx+1)
    print('\n training accuracy '+str(accuracies))

''' Test&Validate your network '''
def test(model, loader,criterion): # you are free to change parameters

    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (data, labels) in enumerate(loader):
            data = data.float() #.cuda()
            labels = labels.float() #.cuda()
            
            #forward only
            y_pred = model(data)

            #loss
            loss = criterion(y_pred, labels)
        
            #assessy_pred = y_pred.flatten()
            y_pred = y_pred.flatten()
            total = labels.size(0)
            correct = (y_pred == labels).sum().item()
            acc = 100 * correct / total
            
            print('validation accuracy '+str(acc))
            return acc
        
        
main()