Diviser un tenseur en torche

Question

Étant donné un tenseur d'entrée de la taille n x 2A x B x C, comment le diviser en deux tenseurs, chacun de taille n x A x B x C? Essentiellement, n est la taille du lot.

Answer 1

Je pense que vous pourriez faire quelque chose comme:

tensor_a = torch.Tensor(n, 2A, B,C) 
-- Initialize tensor_a with the data 

tensor_b = torch.Tensor(n, A, B, C) 
tensor_b = tensor_a[{{},1,{},{}}] 
tensor_c = torch.Tensor(n, A, B, C) 
tensor_c = tensor_a[{{},2,{},{}}]