One more thing, it cannot work with larger size of A_array. I think the purpose to convert to bool to save memory. When I change to int. It shows cuda error out of memory
This is example
import torch
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
B, C, H, W, D=2, 2, 64, 64, 64
input =  torch.randn(B, C, H, W, D)
input = input.to(device)
A_array=np.random.rand(H, W, D) > 0.5
A_array=A_array.astype(int)
print (input)
for ind_batch in range (B):
    input[ind_batch,:, A_array]=0

print (input)


RuntimeError: CUDA error: out of memory