Sorry for the half answer. That seems to be very memory-inefficient. Trying to run the same code on cpu gave me a “Tried to allocate 16GB” error message.
I believe the same output can be obtained with masked_fill. Tried it here and I believe it to be correct.
Generating A_array with torch.rand resolved the memory issue as well for this tensor size, but masked fill ultimately handled somewhat larger tensor sizes without crashing - crashed at (2, 2, 256, 256, 512), which should be a 1GB tensor, a bit much for my notebook.
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 = torch.rand(H, W, D) > 0.5
A_array = A_array.to(device)

input = input.masked_fill(A_array, 0)