python – Is there a Numpy or pyTorch function for this code?

Basically is there a Numpy or PyTorch function that does this:

dims = vp_sa_s.size()
        for i in range(dims(0)):
            for j in range(dims(1)):
                for k in range(dims(2)):
                     #to mimic matlab functionality: vp(mdp_data.sa_s)
                    try:
                        
                        vp_sa_s(i,j,k) = vp(mdp_data('sa_s')(i,j,k))
                        
                    except:
                        print('didnt work with' , mdp_data('sa_s'))

Given that vp_sa_s is size (10,5,5) and each value is a valid index vp i.e in range 0-9. vp is size (10,1) with a bunch of random values.

Matlab do it elegantly and quickly with vp(mdp_data.sa_s) which will form a new (10,5,5) matrix. If all values in mdp_data.sa_s are 1, the result would be a (10,5,5) tensor with each value being the 1st value in vp.

Does a function or method that exists that can achieve this in less than O(N^3) time as the above code is terribly inefficient.

Thanks!