import torch

N = 6
C = 4

data = torch.randn(N, C)
head_mask = torch.tensor([1, 1, 0, 0, 1, 0])
head_index = (head_mask == 1).nonzero().view(-1) # 0, 1, 4
head_index_expand = torch.cat([head_index, torch.tensor([head_mask.numel()]) ] ) # 0,1, 4, 6
head_index_diff = (head_index_expand - torch.roll(head_index_expand, 1))[1:] # 1, 3, 2
splits = torch.split(data, head_index_diff.tolist())

print('ori data')
print(data)
print('after spliting')
print(splits)

# Each item contains a word, using combine all the elements in each item in the list
# you get the words you want.