Solved by eqy in post #2 
              
              
                I think you can use torch.where here: 
import time
import torch

n = 2048
a = torch.randn(2, n)
b = torch.randn(n) > 0
t1 = time.time()
out1 = torch.empty(n)
for i in range(n):
    if b[i]:
        out1[i] = a[1][i]
    else:
        out1[i] = a[0][i]
t2 = time.time()
out2 = torch.where(b, a[1], a[0…