I reduced the running time from 681 seconds to 360 seconds for 50 training epochs. The code was optimized by processing the raw data and storing it as two big tensors in the .pt format.