SessionBatchSampler#

class core.samplers.SessionBatchSampler(sampler, batch_size, *, shuffle_batches=False, generator=None, drop_last=True)[source]#

Bases: torch.utils.data.sampler.Sampler

Group any DatasetIndex sampler’s output into per-session batches.

All indices in a batch share the same session id. The inner sampler controls the order of indices within each session; shuffle_batches controls the order of batches across sessions.

Parameters:

  • sampler (Sampler) – Inner sampler whose __iter__ yields DatasetIndex objects.

  • batch_size (int) – Number of samples per batch.

  • shuffle_batches (bool) – Whether to shuffle the final batch order. Defaults to False.

  • generator (Optional[Generator]) – Generator used when shuffling batches. Defaults to None.

  • drop_last (bool) – Whether to drop the last incomplete batch per session. Defaults to True.

set_epoch(epoch)[source]#

Pass the epoch to the underlying sampler if it supports it.

Return type:

None