What is your current approach?
Not sure if this one looks better than yours:
device = 'cuda'
a = torch.tensor([2, 3, 4], device=device)
length = 5
res = torch.stack([torch.arange(length, device=device) < a_ for a_ in a])