I need to extract elements of a tensor (at present done using a boolean mask), do some operations on them, and replace in the original tensor. The basic operations look like this, where index is a boolean tensor.
new = out[index]

# Do some operations on new here 

rt = out.clone()
rt[index] = new

Is there a faster way of doing this? In the test below, training time is roughly doubled (though depends on how many elements of index are True)
class test(nn.Module):

    def __init__(self):
        super(test, self).__init__()
        self.a = nn.Parameter(torch.ones([100, 200]))
    
    def forward(self, x, index=None):
    
        out = torch.matmul(self.a, x)
    
        if index is not None:
            new = out[index]     
            rt = out.clone()   
            rt[index] = new
        else:
            rt = out
    
        return rt

model = test()
opt = torch.optim.SGD(lr=0.01, params=model.parameters())
loss = torch.nn.MSELoss()

ixs = (torch.arange(100) > 50)

st = time.time()
for _ in range(1000):

    x = torch.ones([200, 400]) + 10
    target = torch.zeros([100, 400])

    out = model(x, ixs)

    l = loss(out, target)

    l.backward()
    opt.step()
    model.zero_grad()

print(time.time() - st)