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?