So the Thomas rule of thumb™ is that you need your tensor sizes to be at least three figures for PyTorch to work well.
The reason for this is that creating tensors carries an overhead and this is likely killing performance in your code when you have lots of <= 3-element tensors. (This is also why torch.jit.script doesn’t really help you here.) It is not entirely clear to me if you compare autograd-enabled PyTorch to NumPy, this would further increase the administrative overhead for each operation. Things will get better if and when you can leverage the JIT fusers as these then operate on “bare metal” internally instead of creating one tensor after another, but these don’t cover all of PyTorch operations and one would need to analyse a bit how to make most of them. The two scripted functions look like they might be very much in scope for them at first sight.
If you can make it into something one can copy and paste to run (so add a self-contained call), it might be interesting to profile it and see if the fusers can help you.
Best regards
Thomas
P.S.: This isn’t the place where you might expect it to be recommended, but I hear that cases similar to this are said to be a nice fit for Julia, which will compile to “bare metal” after differentiation.