Solved by ptrblck in post #2 
              
              
                While a custom CUDA kernel could potentially yield a speedup for the operations, you could start with a pure PyTorch implementation and could thus also use the GPU. 
The numpy methods should also be available in PyTorch and you could just rewrite the code. 
If you see a bottleneck in this translatio…