Solved by ptrblck in post #2 
              
              
                Assuming you are using cudnn, you could add torch.backends.cudnn.benchmark = True at the beginning of your script to use the cudnn heuristics to pick the fastest algorithms for your workload. 
Note that the first iteration for each new input shape will be slow, as cudnn is benchmarking the kernels, …