# Unfolding tensor based on binary map

I have a 4d tensor that I need to perform some operations along the fourth dimension.

I have a 2d binary map for the first two dimensions, where I need to perform the operation only if it’s true.

I found that I can do this to get the sub tensor that I need to perform the operations on:

``````sub_patches=patches[binary_map, :, :]
``````

which gives the desired time speed up.

but this combines the first two dimension, I need to put back where the values were.

Does anybody know how to do that?

Thanks

Can you just assign back the sub tensor to the original at the desired index?

``````t = torch.arange(0, 36).view(2, 2, 3, 3)
binary_map = (torch.arange(0, 4) % 2 == 0).view(2, 2)
sub_t = t[binary_map, :, :]
sub_t = sub_t * 0
t[binary_map, :, :] = sub_t  # just assigning it back
print(t)

Output:
tensor([[[[ 0,  0,  0],
[ 0,  0,  0],
[ 0,  0,  0]],

[[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17]]],

[[[ 0,  0,  0],
[ 0,  0,  0],
[ 0,  0,  0]],

[[27, 28, 29],
[30, 31, 32],
[33, 34, 35]]]])
``````

thanks. I tried this and thought things were put back in different order. It turns out it’s because I was calculating median and the result I was comparing against took the smaller number when there’s even number, while pytorch took the average.