PyTorch: как добавлять тензоры в список с помощью циклов

Question

У меня есть следующий код, который выводит 2 массива в список:

arr1 = np.array([[1.,2,3], [4,5,6], [7,8,9]])

arr_split = np.array_split(arr1,
                           indices_or_sections = 2,
                           axis = 0)

arr_split

Вывод:

[array([[1., 2., 3.],
        [4., 5., 6.]]), array([[7., 8., 9.]])]

Как преобразовать эти 2 массива в тензоры PyTorch и поместить их в список с использованием циклов for (или while), чтобы они выглядели так:

[tensor([[1., 2., 3.],
         [4., 5., 6.]], dtype=torch.float64),
tensor([[7., 8., 9.]], dtype=torch.float64)]

Заранее большое спасибо!

Answer 1

Лучше сначала преобразовать его в тензор, а затем использовать torch.Tensor.split

arr1 = np.array([[1.,2,3], [4,5,6], [7,8,9]])
t_arr1 = torch.from_numpy(arr1)

t_arr1.split(split_size=2)
(tensor([[1., 2., 3.],
        [4., 5., 6.]], dtype=torch.float64), 
 tensor([[7., 8., 9.]], dtype=torch.float64))