Hi guys. I’m trying to do the classification task.
and want to check my code with you guys

is there any part is wrong with pytorch grammar matter
and is there better way to write code? any suggestion

def save_model(model, optimizer, scheduler,idx) :
    model_cpu = model.to('cpu')
    state = {
        'model' : model_cpu.state_dict(),
        'optimizer' : optimizer.state_dict(),
        'scheduler' : scheduler.state_dict()
    }
    if not(os.path.isdir('./saved_model')) : os.mkdir('./saved_model')
    torch.save(state, './saved_model/saved_model.pth')


def train():
    model.train().cuda()
    print('==============================train===============================')
    lst_out = []
    lst_label=[]
    for idx2,(images,labels) in enumerate(train_loader):
        images,labels = torch.tensor(images),torch.tensor(labels)
        images = images.cuda()
        labels = labels.cuda()
        outputs = model(images).cuda()
        loss = criterion(outputs, labels)
        if idx2%args.iter_num ==0: print('train->{}, loss->{}'.format((idx2)/(len(train_loader)),loss))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()



def valid():
    print('==============================valid===============================')
    torch.cuda.empty_cache()
    lst_out = []
    lst_label=[]
    with torch.no_grad():
        model.eval().cuda()
        for idx,( image, label) in enumerate(valid_loader):
            image = image.cuda()
            pred = model(image)
            lst_out += torch.max(pred,1).indices
            label = label.cuda()
            lst_label += label
        f1_scr = f1_score(torch.tensor(lst_out), torch.tensor(lst_label), average='macro')
    return f1_scr



def test() :
    print('==============================test===============================')
    torch.cuda.empty_cache()
    test_loader = data_loader(phase='test', batch_size= args.ts_bs)
    os.system('rm -rf ./test_result.csv')
    f1 = open('./test_result.csv', 'w', encoding ='utf-8-sig', newline='')
    wr = csv.writer(f1)
    lst_out = []
    lst_label = []
    ct = 0
    model = MODEL
    model.load_state_dict(torch.load('./saved_model/saved_model.pth')['model'])
    with torch.no_grad():
        torch.eval().cuda()
        for idx,( image, _) in enumerate(test_loader):
            image = image.cuda()
            pred = model(image)
            pred = pred.argmax(1).cpu().detach().numpy().astype('int')
            lst_out.extend(pred)
            if idx%iter_num:print('idx[',idx,']/',len(test_loader),len(lst_out))
        for idx, item in enumerate(lst_out) :
            wr.writerows([[dic_label[item]]])


if __name__ == "__main__":

    f1_scr = 0
    model = MODEL
    model.cuda()
    criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.001)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5, gamma =0.1)
    best_acc,loss = 0,0

    train_loader = data_loader(phase='train', batch_size=args.bs)
    valid_loader = data_loader(phase='validation',batch_size=8)
    print('main')

    for i in range(args.epoch):
        train()
        f1_result = valid()
        if f1_result > f1_scr:
            f1_scr = f1_result
            save_model(model,optimizer,scheduler,idx)
    #if args.status =='train':train()