PyTorch torch.where función
Pytorch torch manual de referencia
torch.whereEs una función en PyTorch que se utiliza para devolver elementos según una condición.
Definición de la función
torch.where(condition, input, other)
Ejemplo de uso
Ejemplo
import torch
condition = torch.tensor([[True, False], [False, True]])
x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[10, 20], [30, 40]])
result = torch.where(condition, x, y)
print(result)
condition = torch.tensor([[True, False], [False, True]])
x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[10, 20], [30, 40]])
result = torch.where(condition, x, y)
print(result)
El resultado de salida es:
tensor([[ 1, 20],
[30, 4]])
Otras extensiones