CODEWITHSUNDEEP
pytorch
Minh-Long Luu
Minh-Long Luu

Reputation: 2741

PyTorch: apply mask with different shape

I have a tensor of shape (60, 3, 32, 32) and a boolean mask of shape (60, 32, 32). I want to apply this mask to the tensor. The output tensor should have shape (60, 3, 32, 32), and values are kept if the mask is 1, else 0. How can I do that fast?

Upvotes: 1

Views: 759

Answers (1)

GoodDeeds
GoodDeeds

Reputation: 8527

Let t be the tensor and m be the mask. You can use:

t * m.unsqueeze(1)

Upvotes: 3

Related Questions