This one will definitely be more efficient as the for loop is only reading an attribute on the Tensor and not indexing a Tensor.
Also the split is equivalent to what I meant with narrow and will be fast.
If you want to sum all these Tensors you will have to add a new for-loop I’m afraid.
If you work on CPU it shouldn’t be too bad though.
Note that we are currently working on adding NestedTensors to pytorch: https://github.com/pytorch/nestedtensor 19
This will allow you to handle different-sized Tensors as a single Tensor.