In the middle of my model, there is some for loops that I must do, there is no replace for it, so to make it faster I will use numba.njit, but it works for only numpy arrays and also if I detach the tensor the gradients won’t be recorded. So is there a solution for that, and also I heard of torch.jit.trace, will it do the same as numba.njit and record the gradients.