@DeepC004 sorry for delay and I just checked the runtime and this example adds +10 minutes of CI runtime. Let's try to reduce that

Hi, one way to reduce the training time can be increasing the batch size. I am unaware of the memory limits of the remote host used for CI, hence one way to try is to keep increasing the train and test batch size in powers of 2. Also, training time is the bottleneck hence increasing the num_workers while dataloading will not make difference, nevertheless I will play around to see if there are any improvements in runtime.