Filter torch tensor of multidimensional array

data
1

I have a model that responds with a tensor object and I want to filter it so that I can get a new tensor.
My goal is to filter by the 6th argument so that I can get only the values that are zero (0).
How would I do this efficiently?

I have an example of such an object. I don’t want to create a new object every time because it takes a while to create it as a cuda tensor.

import torch
myobj = torch.tensor([
   [8.32500e+02, 1.72969e+02, 9.27000e+02, 3.81375e+02, 6.51855e-01, 0.00000e+01], 
   [5.01750e+02, 1.21500e+02, 8.13000e+02, 3.91875e+02, 2.75879e-01, 6.10000e+01], 
   [7.90500e+02, 2.60625e+02, 8.53500e+02, 3.55875e+02, 2.71729e-01, 0.00000e+01], 
   [0.00000e+00, 2.12812e+01, 1.26562e+02, 3.68250e+02, 2.68066e-01, 7.20000e+01]
])
myobj = myobj.to('cuda')
2

If I understood your question correctly, you want to filter myobj and in your example only return the first and third rows. You can do this with

filtered_myobj = myobj[myobj[:, 5] == 0]
3

Thanks! It works perfectly!
I needed the ones that are equal with 0 so I used it like that:

filtered_myobj = myobj[myobj[:, 5] == 0]
1 Like
4

Of course! Sorry about that. I’ve updated my answer so no one else who reads this gets confused by the accepted solution.