Torch equivalent of np.where

1

Hi,

I am working with 4-dimensional tensors, and I am using for loop to check the conditional flow, I wanted to know how can I implement the below code efficiently

## object_temp Shape -> (24, 7, 7, 30)
## Object Shape -> (24, 7, 7)
## target Shape -> ( 24, 7, 7, 30)

        for i in range(object.shape[0]):
          for j in range(object.shape[1]):
            for k in range(object.shape[2]):
              if ((target[i,j,k,4] >0) | (target[i,j,k,9] >0)):
                object[i,j,k] = 1
                object_temp[i,j,k,:] = 1
2
3

It does not take in input of 1 and 0,

I want to create a tensor of the same dimensions as the original tensor target.

I can only create one less dimension tensor using torch.where. How can I also get the 4 dimension?

# Wanted object_temp Shape -> ( 24, 7, 7, 30)
#target Shape -> ( 24, 7, 7, 30)
4

I don’t understand at all why you need to do a nested loop. You are not permuting indices and you are always taking i,j,k in an ordered way. You can just do something like

mask = target[...,4]>0 |  target[...,9]
object[mask]=1
object[mask.expand_as(object_temp)]=1
1 Like