If I understand the issue correctly at the moment, you are seeing a time of 0.117 s/epoch using random data on the GPU and 1.22 s/epoch if you use your real data?
We are aware of denormal values, which might slow down the execution on the CPU, but this shouldn’t be the case on the GPU.
Could you nevertheless set torch.set_flush_denormal(True) and time it again?