Well, for best perf, you want to move your function into the inner loop in cpp or cuda code.
Which you cannot do using a python function.
Also remember that NestedTensors are built on top of pytorch. So if you can do your code with NestedTensor, you can do it without