torch.arrange

1. 세로로 긴 2차원 행렬

 positions_list = torch.arange(0, max_len, dtype=torch.float).view(-1, 1) # 0, 1, 2, 3, 4, 5

 

tensor([[0.],

            [1.],

            [2.],

            [3.],

            [4.],

            [5.]])

 

2. step 간격의 배열

import torch

dim_model = 6

positions_list =  torch.arange(0, dim_model, 2)

print(positions_list)

 

result

tensor([0, 2, 4])