Speed up CNN pytorch

Answered here.
Have you profiled the code and checked, where the bottleneck are?