I have a tensor, A of shape [X,Y,Z] and an indexing tensor, B, of [X,Y,1]. I would like to index A with B, and replace those elements of A with 0.
B is the output of a neural network, with the desired result that it is either true or false for each Y.
r1 = -10
r2 = 10  #some ranges

X = 20
Y=100
Z=5

activation = nn.Tanh() #an activation function
 
#A is an arbitrary random tensor [0,1]
A = torch.FloatTensor(X,Y,Z).uniform_(0, 1)

#B is an arbitrary random tensor [-1,1]
B = activation(torch.FloatTensor(X,Y,1).uniform_(r1, r2))

#idx are where B is less than 0
idx = (B<0)

#Some function to mask A with B such that where B is True, A=0
A = foo(A,B)

What is the correct function, foo, required such that A is both masked with zero and differentiable?