CODEWITHSUNDEEP
pythonnumpybooleanpytorchtensor
Maximilian
Maximilian

Reputation: 1375

Elegant way to compare to torch.FloatTensor on GPU

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?

Upvotes: 5

Views: 2664

Answers (2)

timbmg
timbmg

Reputation: 3328

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.

Upvotes: 4

kmario23
kmario23

Reputation: 61345

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

Upvotes: 1

Related Questions