I’m wondering if the reason why it is slow is because it’s mostly implemented in python instead of the C/C++ backend. I’ve opened a feature request to implement it in the C/C++ backend here: https://github.com/pytorch/pytorch/issues/5155 6