Pytorch: torch.split против индексации в отношении производительности

Question

Я хотел бы выполнять партии тензора в PyTorch отдельно. Что быстрее: indexing или torch.split, чтобы получить элементы пакета и обработать их позже. Т.е.:

# Number of batches
n = x.shape[0]  

# Option 1
x_groups = x.split(x, n)

# Option 2
x_groups = [x[i] for i in range(n)]

# Separate operations
for xi in x_groups: 
  # do something

Предположим, что у обоих методов есть производные.

Спасибо