Hi,
I think the point was that for such a small model, the overhead of running the python code is actually going to be noticeable compared to something that runs on pure c++.
The printing every 10 batch can slow down the process quite a lot as well.
Finally, even though it should be minimal here, you need to properly synchronize the gpu with torch.cuda.synchronize() to get accurate time measurements of gpu compute.