How to conditionally construct a tensor from two other tensors in PyTorch on the GPU?
An example:
import torch

pred = torch.tensor([1,2,1,0,0], device='cuda:0')
correct = torch.tensor([1,0,1,1,0], device='cuda:0')
assigned = torch.tensor([1,2,2,1,0], device='cuda:0')

I want result = tensor([1,2,1,1,0], device='cuda:0').

Basically, when pred is same as correct then correct else assigned.

Moreover, I want to exclude this computation from the gradient computation.


Is there a way to do this without iterating over the tensors?

torch.where does exactly what you're looking for:
import torch

pred = torch.tensor([1,2,1,0,0], device='cuda:0')
correct = torch.tensor([1,0,1,1,0], device='cuda:0')
assigned = torch.tensor([1,2,2,1,0], device='cuda:0')

result = torch.where(pred == correct, correct, assigned)

print(result)
# &gt;&gt;&gt; tensor([1, 2, 1, 1, 0], device='cuda:0')

Since none of these tensors have requires_grad=True, nothing needs to be done to avoid gradient computation. Otherwise, you can do something like this:
import torch

pred = torch.tensor([1.,2.,1.,0.,0.], device='cuda:0')
correct = torch.tensor([1.,0.,1.,1.,0.], device='cuda:0', requires_grad=True)
assigned = torch.tensor([1.,2.,2.,1.,0.], device='cuda:0', requires_grad=True)

with torch.no_grad():
    result = torch.where(pred == correct, correct, assigned)

print(result)
# &gt;&gt;&gt; tensor([1, 2, 1, 1, 0], device='cuda:0')

If you do not use torch.no_grad(), you'll have:
result = torch.where(pred == correct, correct, assigned)
print(result)
# &gt;&gt;&gt; tensor([1., 2., 1., 1., 0.], device='cuda:0', grad_fn=&lt;SWhereBackward&gt;)

which, then, can be detached from the computational graph by using:
result = result.detach()
print(result)
# &gt;&gt;&gt; tensor([1., 2., 1., 1., 0.], device='cuda:0')

