Alright, I measured the forward pass time using torch.cuda.synchronize() this time :
            mean_time = 0
            for _ in tqdm.tqdm(range(1000)):
                torch.cuda.synchronize()
                start_t = time.time()
                mod(model_input)
                torch.cuda.synchronize()
                end_t = time.time()
                mean_time += (end_t - start_t) / 1000

The results are not as bad as before indeed, but the time still goes from 5ms to 12ms between batch sizes 1 and 2, the difference is still big enough to make me wonder if this is to be expected. Do you think this issue is related to my code or cuda related ?