How to apply mask to image tensors in PyTorch?
Applying mask with NumPy or OpenCV is a relatively straightforward process. However, if I need to use masked image in loss calculations of my optimization algorithm, I need to employ exclusively PyTorch, as doing otherwise interferes with gradient computations.
Assuming that I have an image tensor [1, 512, 512, 3] (batch, height, width, channels) and a mask tensor [1, 20, 512, 512] (batch, channels, height, width) where every channel corresponds to one of 20 segmentation classes, I want to get a masked image tensor that fills every pixel with black (0, 0, 0), except for those belonging to one or more specified segmentation classes.
Here is how it can be done with numpy:
import numpy as np
import torch

# Create dummy image and mask
image_tensor = torch.randn([1, 512, 512, 3])
mask_tensor = torch.randn([1, 20, 512, 512])

# Apply argmax to mask
mask_tensor =  torch.max(mask_tensor, 1)[1] # -&gt; 1, 512, 512

# Define mask function
def selective_mask(image_src, mask, dims=[]):
    h, w = mask.shape
    background = np.zeros([h, w, 3], dtype=np.uint8)
    for j_, j in enumerate(mask[:, :]):
        for k_, k in enumerate(j):
            if k in dims:
                background[j_, k_] = image_src[j_, k_]
    output = background
    return output

# Convert tensors to numpy:
image = image_tensor.squeeze(0).cpu().numpy()
mask = mask_tensor.squeeze(0).cpu().nmpy()

# Apply mask function for several classes
image_masked = selective_mask(image, mask, dims=[5, 6, 8])

How should my code be changed to bring it in line with the PyTorch requirements?

First of all, the definition of the function selective_mask is far for what You may call 'straightforward'. The key point in using numpy (and torch, which is designed to be mostly compatible) is to take advantage of the vectorization of operations and to avoid using loops, which are not parallelizable.
If You rewrite the said function in this manner:
def selective_mask(image_src, mask, channels=[]):
    mask = mask[np.array(channels).astype(int)]
    return np.sign(np.sum(mask, axis=0), dtype=image_src.dtype) * image_src

it will turn out that You can actually do the same with pytorch tensors (here no need to squeeze the batch (first) dimension):
def selective_mask_t(image_src, mask, channels=[]):
    mask = mask[:, torch.tensor(channels).long()]
    mask = torch.sgn(torch.sum(mask, dim=1)).to(dtype=image_src.dtype).unsqueeze(-1)
    return mask * image_src

Also, You probably want to produce the mask itself this way:
(BTW using a combination of max and sgn here should actually work faster than setting elements indexed by argmax)
# Create dummy image and mask
image_tensor = torch.randn([1, 512, 512, 3])
mask_tensor = torch.randn([1, 20, 512, 512])

# Discreticize the mask (set to one in the channel with the highest value) -&gt; 1, 20, 512, 512
mask_tensor = torch.sgn(mask_tensor - torch.max(mask_tensor, 1)[0].unsqueeze(1)) + 1.  

Then it should work just fine:
print(selective_mask_t(image_tensor, mask_tensor, [5, 6, 8]))

