python - 与 GPU 上的 torch.FloatTensor 进行比较的优雅方式

Question

我尝试像这样比较位于 GPU 上的两个 torch.FloatTensor(只有一个条目):

if (FloatTensor_A > FloatTensor_B): do something

问题是，(FloatTensor_A > FloatTensor_B) 返回 ByteTensor。有没有一种方法可以在这两个标量 FloatTensors 之间进行 boolean 比较，而无需在 CPU 上加载张量并将它们转换回 numpy 或传统 float ？

Answer 1

PyTorch 中的比较操作返回 ByteTensors(参见 docs)。为了将结果转换回 float 据类型，您可以对结果调用 .float()。例如:

(t1 > t2).float()

(t1 > t2) 将返回一个 ByteTensor。

操作的输入必须在同一内存上(CPU 或 GPU)。返回结果将在同一内存中。当然，任何Tensor都可以通过在其上调用.cpu()或.cuda()移动到各自的内存。