THank you albabD
but did you visit the github i provide?
I try to use the sam optimizer from what they provide
so can you suggest some code modification?
I changed it like this but I guess this is not what they are asking

def train_model_sam(model, epoch,train_loader,val_loader,optimizer,criterion,scheduler):
model.train()
losses = AverageMeter()
accs = AverageMeter()

tk = tqdm(train_loader, total=len(train_loader), position=0, leave=True)
for idx, (imgs, labels) in enumerate(tk):
    imgs_train, labels_train = imgs.cuda(), labels.cuda().long()
    output_train = model(imgs_train)

    # first forward-backward pass
    #optimizer.zero_grad() 
    loss = criterion(output_train, labels_train)  # use this loss for any training statistics
    loss.backward(retain_graph = True)
    optimizer.first_step(zero_grad=True)
    

    # second forward-backward pass
    #optimizer.zero_grad() 
    output_train = model(imgs_train)
    loss=criterion(output_train, labels_train)
    loss.backward()
    optimizer.second_step(zero_grad=True)

    #loss = criterion(output_train, labels_train)
    #optimizer.zero_grad() 
    #loss.backward()
    #optimizer.step() 
    
    accs.update((output_train.argmax(1)==labels_train).sum().item()/imgs_train.size(0),imgs_train.size(0))
    losses.update(loss.item(), imgs_train.size(0))

    tk.set_postfix(loss=losses.avg,acc=accs.avg)
    
return losses.avg