Thank you for your help, I just tried using static input, I use the code below for the training loop
dummy_set = torch.rand(940, 3, 150, 150)
dummy_set.requires_grad=True
dummy_lbl = torch.randint(low=0, high=5, size=(940,))

import time
torch.backends.cudnn.benchmark =True 
torch.cuda.synchronize()
s = time.time()
 
model.train()
for i in range(epoch):
    total_loss = 0
    total_sample = 0
    total_correct = 0

    bst = 0  # batch start
    bsz = 64 # batch size
    ben = bst*bsz+bsz # batch end

    while ben!=len(dummy_set):
        image, label = dummy_set[bst*bsz:ben], dummy_lbl[bst*bsz:ben]
        image = image.to('cuda')
        label = label.to('cuda')
        out = model(image) 
        loss = criterion(out, label)
        optimizer.zero_grad()
        loss.backward() 
        optimizer.step() 
        total_loss += loss 
        total_sample += len(label)
        total_correct += torch.sum(torch.max(out,1)[1]==label)*1.0

        bst += 1 # increase batch
        ben = min(bst*bsz+bsz, len(dummy_set))

    print(f"epoch {i} loss:{total_loss/total_sample}  acc:{total_correct/total_sample}") 

torch.cuda.synchronize()
e = time.time()

and the result, the gap now is smaller but K80 is still faster ( I run it several times):

K80 : ~56s
P6000 : ~60s