Expand the 'start' and 'end' parameter of torch arange

1

Hello all~

With torch.arange(start=1, end=4) I will have a tensor [1, 2, 3]
Now if I want to replace the ‘start’ parameter’s value with a tensor, how can I do that without doing the loop?
for example if I have a=[[1], [2]]
If I write torch.arange(start=a, end=a+3)
I want to obtain
[[1, 2, 3],
[2, 3, 4]]

Is it possible to do that without doing the loop?

2

Hi Duy!

You can add the desired start values to the result of arange(), and do
so on a tensor basis:

>>> import torch
>>> torch.__version__
'1.9.0'
>>> start_values = torch.tensor ([1, 2, 5, 7])
>>> torch.arange (3).repeat (4, 1) + start_values.unsqueeze (1)
tensor([[1, 2, 3],
        [2, 3, 4],
        [5, 6, 7],
        [7, 8, 9]])

Best.

K. Frank

1 Like
3

Great solution! Thank you very much Frank!