Fastest way to split tensors in groups of two and for each group, randomly assign one to 0 and the other to 1
I have a big tensor with size (128, 64, 1, 1). For every two small tensors of size (1, 1) in that big tensor, I want to randomly assign one to be 1 and the other to be 0.

My following code works but it's awfully slow. 

# a is a big tensor of size (128, 64, 1, 1)
for i in range(a.size()[0]):
    j = 0
    while j &lt; a.size()[1] - 1:
        r = int(torch.randint(0, 2, (1,)))
        a[i][j + r] = 1
        a[i][j + 1 - r] = 0
        j += 2



Is there any way to avoid for loops and do every assignment in parallel? Thank you.

I think this might be faster :

import torch
x = torch.randn(128, 32, 2, 1, 1)
y = x.max(dim=2, keepdim=True)[0]
z = (y-x) &gt; 0
r = torch.flatten(z, 1, 2)


The idea is that you split your 64 dimension into 2, pick random values, and chose the one that is minimal to be 1 and the other to be 0, then you recombine your 32*2 into a 64 size.
