I have an input tensor of shape: data (x,y,z). I have a binary mask tensor of shape: mask (x,y).
When I do data[mask > 0] I obtain a new tensor of shape (q,z) where q is the number of ones in the mask tensor.
I would instead want to get a tensor of original shape (x,y,z) but for the values for which we have zeros in mask being eliminated and instead being padded at the end with 0 in data tensor so we keep the original shape ( the reason it doesn't do that now is because we would have variable length across second dimension).
Of course this can be easily done in python with basic matrix operations, but is there an efficient tensor-way to do it in pytorch?
Example (imagine a,b,c...are 1D tesnors):
data
[[a,b,c],
[d,e,f]]
mask
[[0,1,0],
[1,0,1]]
Ideal output:
[[b,junk,junk],
[d,f, junk]]
Here, the missing stuff is padded with some "junk" to keep the original shape.