Grain (JAX) — эквивалент pyTorch `collect_fn` для пакетов. ⇐ Python

[Anonymous]

Я определил класс набора данных с помощью __len__ и __getitem__, который возвращает кортеж значений. Я могу использовать `grain.transforms.Batch` для составления пакетов, но как указать, как каждый элемент объединяется в пакет?
Создаю это на будущее...

Подробнее здесь: https://stackoverflow.com/questions/798 ... or-batches