Given a batch of error maps, I would like to mask out error values based on the max of each error map in the batch.
Using torch.max currently gives the single maximum value in the whole batch of error maps. Masking using this single value would result in inconsistent masks with respect to each error map.
Example Input: Loss Map shape: (8, 256, 640) (Batch, Height, Width)
Return: (8, 256, 640) with masked out pixels.
Current Implementation:
bound = torch.max(loss[i]) * 0.9          # Find threshold
mask = torch.zeros_like(loss[i])          # create mask
mask[loss[i] < bound] = 1.0                # set the values i want to keep to 1
mask[loss[i] >= bound] = 0.0             # set others to 0
loss[i] *= mask                                  # Mask!

I am looking for an efficient way to do this in PyTorch instead of a naive approach of using a for loop for each image in a batch.
Thanks!