Your code looks generally fine beside some small things I would change.


outputs = model(images).cuda(): the cuda() call is not necessary, since the output should already be on the GPU (or are you transferring it to the CPU inside the model?) Also, to write device-agnostic code, you could use:

device = 'cuda:0'`
images = images.to(device)
labels = labels.to(device)
...



torch.cuda.empty_cache(): Clearing the cache is usually not necessary and will slow down your code. You might want to use it, if you expect other applications to use GPU memory and thus want PyTorch to release it.


model.eval().cuda(): unnecessary cuda() operation


NIT: lst_out += torch.max(pred,1).indices vs. pred = pred.argmax(1).cpu().detach().numpy().astype('int'): you are getting the predictions in both lines of code, but use quite a different coding style. Not an issue at all, so just a minor remark.


model = MODEL: this looks a bit fishy and I guess MODEL might just be a placeholder for this post?


NIT: step_size = 5, gamma =0.1 and best_acc,loss = 0,0: your spacing is inconsistent. You could use a linter to clean it up.


train() and valid() method use global variables. Personally I would pass all variables as arguments as I’ve ran into issues, where I’ve changed the global variable name in the script and the method was still using an “old” variable.