vectorize pytorch tensor indexing
I have a batch of images img_batch, size [8,3,32,32], and I want to manipulate each image by setting randomly selected pixels to zero. I can do this using a for loop over each image but I'm not sure how to vectorize it so I'm not processing only one image at a time.  This is my code using loops.
batch_size = 8
prct0 = 0.1
noise = torch.tensor([9, 14, 5, 7, 6, 14, 1, 3])
comb_img = []

for ind in range(batch_size):

    img = img_batch[ind]
    c, h, w = img.shape          
    prct = 1 - (1 - prct0)**noise[ind].item()
    idx = random.sample(range(h*w), int(prct*h*w)  )
    img_noised = img.clone()
    img_noised.view(c,1,-1)[:,0,idx] = 0 
    comb_img.append(img_noised)

comb_img = torch.stack(comb_img) # output is comb_img [8,3,32,32]

I'm new to pytorch and if you see any other improvements, please share.

First note: Do you need to use noise? It will be a lot easier if you treat all images the same and don't have a different set number of pixels to set to 0.
However, you can do it this way, but you still need a small for loop (in the list comprehension).
#don't want RGB masking, want the whole pixel
rng = torch.rand(*img_batch[:,0:1].shape) 
#create binary mask
mask = torch.stack([rng[i] &lt;= 1-(1-prct0)**noise[i] for i in range(batch_size)]) 
img_batch_masked = img_batch.clone()
#broadcast mask to 3 RGB channels
img_batch_masked[mask.tile([1,3,1,1])] = 0

You can check that the mask is set correctly by summing mask across the last 3 dims, and seeing if it matches your target percentage:
In [5]:     print(mask.sum([1,2,3])/(mask.shape[2] * mask.shape[3]))
tensor([0.6058, 0.7716, 0.4195, 0.5162, 0.4739, 0.7702, 0.1012, 0.2684])

In [6]:     print(1-(1-prct0)**noise)
tensor([0.6126, 0.7712, 0.4095, 0.5217, 0.4686, 0.7712, 0.1000, 0.2710])

