Elegant way to compare to torch.FloatTensor on GPU

user_4689412 0
Viewed 2094

I try to compare two torch.FloatTensor (with only one entry) lying on GPU like this:

if (FloatTensor_A > FloatTensor_B): do something

The problem is, that (FloatTensor_A > FloatTensor_B) gives ByteTensor back. Is there a way to do boolean comparison between these two scalar FloatTensors, without loading the tensors on CPU and converting them back to numpy or conventional floats?

tensor pytorch boolean numpy python
2 Answers

The comparison operations in PyTorch return ByteTensors (see docs). In order to convert your result back to a float datatype, you can call .float() on your result. For example:

(t1 > t2).float()

(t1 > t2) will return a ByteTensor.

The inputs to the operation must be on the same memory (either CPU or GPU). The return result will be on the same memory. Of course, any Tensor can be moved to the respective memory by callin .cpu() or .cuda() on it.

user_41390240

Yes, it's easy and simple.

Example:

In [24]: import os

# select `GPU 0` for the whole session
In [25]: os.environ['CUDA_VISIBLE_DEVICES'] = '0'

# required `data type` (for GPU) 
In [26]: dtype = torch.cuda.FloatTensor

# define `x` & `y` directly on GPU
In [27]: x = torch.randn(100, 100).type(dtype)
In [28]: y = torch.randn(100, 100).type(dtype)

# stay on GPU with desired `dtype`
In [31]: x.gt(y).type(dtype)
Out[31]: 

    0     1     1  ...      0     0     0
    1     0     0  ...      1     0     1
    1     1     1  ...      0     0     0
       ...          ⋱          ...       
    1     1     1  ...      0     0     0
    0     1     1  ...      1     1     1
    1     0     1  ...      1     0     1
[torch.cuda.FloatTensor of size 100x100 (GPU 0)]


# sanity check :)
In [33]: x.gt(y).type(dtype).is_cuda
Out[33]: True
user_29560660
Related