Определенный пользователем collate_fn в DataLoader изменяет структуру данных

Question

У меня есть собственная функция сортировки, которая изменяет форму тензоров в партии (заполнение переменной длины и т. Д.).Когда я запускаю итерацию дважды (или несколько эпох), она выдает ошибку.После исследования это потому, что вход для второй итерации был выходом из первой!Таким образом, он проходит через функцию сортировки дважды.Как этого избежать?

train_dataset = DataLoader(dataset=train_set, collate_fn=ER_Collate,
                               batch_size=BATCH_SIZE, shuffle=True)
for i, train_ in enumerate(train_dataset):
   pass
for i, train_ in enumerate(train_dataset)://error thrown
   pass