Conditional indexing in Pytorch C++

I have tried unsuccessfully for several days to implement a simple indexing task for
Pytorch in C++. Here is a trivial example. In MATLAB, I use:

x=randn(5,5); % create 5x5 random matrix
z=exp(x); % create other matrices of same type and shape as x
ii=find(x>10); % indexes to valuesgreater than 10
z(ii) = 1./x(ii); % get and set values at specified indexes,
x(ii) = z(ii); % while re-using indexes “ii” to save time.

In Python:

ii = np.where(x > 10)[0]
z[ii] = 1/x[ii]
x[ii] = z[ii]

Here is my C++ implementation that does not work:

torch::Tensor test_fn(torch::Tensor x) {
    using namespace torch::indexing;
    auto z =  x.exp();
    auto mask = x>10;
    // This works, but does not allow me to re-use the index:
    //     z.index_put_({mask}, 1/x.index({mask}));

    // auto ii = mask.nonzero();  // tried this
    // ii = at::nonzero(mask);   // tried this too
    auto ii = torch::nonzero(mask).squeeze();
    // z.index_put_({ii}, 1/x.index({ii}));  // tried this
    z.index_put_({ii}, 1/x.index_select(0,torch::tensor(ii)));
    return z;
````Preformatted text`

Can someone help me with this? Note that it is important that I have
the required indexes, so I can re-use them. The number of elements satisfying
the condition are relatively few, and the operations required on these few indexes
is large (not like in my trivial example). So it is important to do the operations
only on those indexes and re-use the indexes.

Could you describe the issue in more detail and what exactly is failing?
Based on your comments it seems:

// This works, but does not allow me to re-use the index:

you are unable to re-use the indices? Do you see an error?

Wheh I use the line:
z.index_put_({mask}, 1/x.index({mask}));
I get the desired result, but I do not have a set of indexes, analogous to “ii” in the MATLAB code I provided.

Would torch::where work as seen in the corresponding numpy example?

Torch where and numpy where are fundamentally different. torch::where is like an if/else, but numpy.where returns indexes to the elements of a boolean vector that are true. Although I have working code using index_put_ , I still would like to know how to implement MATLAB “find”, or else Python “where” in torch or C++, whereby I gt a list of the elements that are nonzero or true.