Table of Contents
Shortcuts

# 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:

  out i={input i if condition i other i otherwise\text{out}_i = \begin{cases} \text{input}_i & \text{if } \text{condition}_i \\ \text{other}_i & \text{otherwise} \\ \end{cases} out i​={input i​other i​​if condition i​otherwise​

  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:

    A tensor of shape equal to the broadcasted shape of `condition`, `input`, `other`

  Return type:

    Tensor

  Example:

      >>> x = torch.randn(3, 2)
      >>> y = torch.ones(3, 2)
      >>> x
      tensor([[-0.4620,  0.3139],
              [ 0.3898, -0.7197],
              [ 0.0478, -0.1657]])
      >>> torch.where(x > 0, 1.0, 0.0)
      tensor([[0., 1.],
              [1., 0.],
              [1., 0.]])
      >>> torch.where(x > 0, x, y)
      tensor([[ 1.0000,  0.3139],
              [ 0.3898,  1.0000],
              [ 0.0478,  1.0000]])
      >>> x = torch.randn(2, 2, dtype=torch.double)
      >>> x
      tensor([[ 1.0779,  0.0383],
              [-0.8785, -1.1089]], dtype=torch.float64)
      >>> torch.where(x > 0, x, 0.)
      tensor([[1.0779, 0.0383],
              [0.0000, 0.0000]], dtype=torch.float64)

  torch.where(condition) → tuple of LongTensor

  `torch.where(condition)` is identical to `torch.nonzero(condition, as_tuple=True)`.

  Note

  See also `torch.nonzero()`.

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.