Scalar type in torch.where?

user_11939660 0
Viewed 1813

torch.where documentation states that x and y can be either a tensor or a scalar. However, it doesn't seem to support float32 scalar.

import torch

x = torch.randn(3, 2)  # x is of type torch.float32

torch.where(x>0, 0, x)  # RuntimeError: expected scalar type long long but found float 
# torch.where(x>0, 0.0, x)  # RuntimeError: expected scalar type double but found float

My question is how to use float32 scalar?

pytorch
2 Answers

It's not that torch doesn't support float32. It's your system doesn't provide an easy way to specify 0 as a float32. As stated in the errors, 0 is interpreted as a long long C type, i.e. int64, while 0.0 is interpreted as double C type, i.e. float64.

I guess you need to cast 0 to the same dtype with that of x:

torch.where(x>0.0, torch.tensor(0, dtype=x.dtype), x)
user_42384080

There is a known issue issue in torch 1.9.0 (and potentially earlier versions) where type promotion does not occur. See issue #64161. This is listed as a high priority issue with initial patches written.

user_27808030
Related