Pytorch speedup