torch.where
torch.where(condition, input, other, *, out=None) -> Tensor
Return a tensor of elements selected from either "input" or
   "other", depending on "condition".
The operation is defined as:
  \text{out}_i = \begin{cases}     \text{input}_i & \text{if }
  \text{condition}_i \\     \text{other}_i & \text{otherwise} \\
  \end{cases}

Note:
 The tensors "condition", "input", "other" must be broadcastable.

Parameters:
      * condition (BoolTensor) -- When True (nonzero), yield
        input, otherwise yield other
  * **input** (*Tensor** or **Scalar*) -- value (if "input" is a
    scalar) or values selected at indices where "condition" is
    "True"

  * **other** (*Tensor** or **Scalar*) -- value (if "other" is a
    scalar) or values selected at indices where "condition" is
    "False"

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Returns: