How can I change part of a PyTorch tensor based on the values of another tensor?

Question

This question may not be clear, so please ask for clarification in the comments and I will expand.

I have the following tensors of the following shape:

mask.size() == torch.Size([1, 400])
clean_input_spectrogram.size() == torch.Size([1, 400, 161])
output.size() == torch.Size([1, 400, 161])

mask is comprised only of 0 and 1. Since it's a mask, I want to set the elements of output equal to clean_input_spectrogram where that relevant mask value is 1.

How would I do that?

Answer 1

Basic usage of a mask:

a = torch.zeros([2,])     # tensor([0., 0.])
b = torch.ones([2,]) * 7  # tensor([7., 7.])

m = torch.tensor([True, False])

a[m] = b[m]               # tensor([7., 0.])

In this case, you have an extra dimension in a and b, but not in mask. This is no problem and will be handled automatically due to broadcasting. Example:

a = torch.zeros([2, 161])
b = torch.ones([2, 161]) * 7

m = torch.tensor([True, False])

a[m] = b[m]
print(a)