Batched torch.arange

1

I want to create a tensor that looks like this:

[ [0, 1, 2, 3, ...],
  [0, 1, 2, 3, ...],
  ...               ]

I know I can do this with a loop and torch.arange:

tc.stack([tc.arange(0, L) for i in range(B)])

However, this method can be inefficient because of big python loops. I’m wondering is there a way to achieve this without loops?

Thanks

2

This should work,

temp = torch.ones(B, L)
result = torch.arange(0, L) * temp
1 Like
3

If you just need a view, then the following should work:

torch.arange(L).expand(B, -1)
1 Like