# Splitting torch tensor by flags

Working with PyTorch tensors, I need to split the batch of items by its flags, so the items in x_batch_one and x_batch_two are lists of pairs that have the same flags.

def pairs_by_flags(self, x_batch, y_batch, max_flag):
#split x_batch on pairs

sizes = (x_batch.size(0)//2, *x_batch.size()[1:] )

items_pos = torch.zeros((max_flag,), dtype = torch.long, device = x_batch.device)

x_batch_one = torch.zeros(sizes, device = x_batch.device)
x_batch_one_pos = 0
x_batch_two = torch.zeros(sizes, device = x_batch.device)

for index, flag in enumerate(y_batch):
x_item = x_batch[index]
i_flag = items_pos[flag]
if i_flag:
x_batch_two[i_flag-1] = x_item
items_pos[flag] = 0
else:
x_batch_one[x_batch_one_pos] = x_item
x_batch_one_pos+=1
items_pos[flag] = x_batch_one_pos

return x_batch_one, x_batch_two


Here is a simple example with single digits:

from torch_helper.myt import thelp
import torch

arr = torch.tensor([2,7,6,2,7,3,2,8,8,6,2,3])

arr1, arr2 =   thelp.pairs_by_flags(arr, arr, 9)

print(arr1)

print(arr2)


Output:

tensor([2., 7., 6., 3., 2., 8.])
tensor([2., 7., 6., 3., 2., 8.])


But after making this naive solution I wonder if it could be improved, perhaps eliminating the loop using some tensor [:,,,:,,,:] magic.