Hi,
The problem is most likely that you’re hitting denormal numbers (float numbers too close to zero are very expensive to work with).
Does adding a torch.set_flush_denormal(True) 14 before the op solves the issue?
This will reduce the precision of the operations for very very small numbers but should remove the slowdown.
EDIT: cross post from github here 18.