torchcell.loader.dense_padding_data_loader.DensePaddingCollater

class DensePaddingCollater(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, follow_batch: list[str] | None = None, exclude_keys: list[str] | None = None)

Bases: object

Collate dataset objects into a mini-batch using dense padding.

Pads data objects from a torch_geometric.data.Dataset to a mini-batch along a new dimension. Data objects can be either of type ~torch_geometric.data.Data or ~torch_geometric.data.HeteroData.

collate_fn(batch: list[Any]) → Any

Collate a batch, fetching from disk first for OnDiskDataset.

Collates a python list of data objects to the internal storage format of torch_geometric.data.DataLoader.

Parameters:

batch – A python list of data objects.

Returns:

A mini-batch of data objects.