loader

DataLoader class for machine learning.

class OFLoader(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, batch_size: int = 1, shuffle: bool = False, follow_batch: List[str] | None = ('coeffs', 'atomic_numbers'), exclude_keys: List[str] | None = None, list_keys: List[str] | None = None, **kwargs)

Data loader for OF-DFT data.

Thin wrapper around torch_geometric.loader.DataLoader, setting “follow_batch” to [“coeffs”] by default, to facilitate splitting of all basis-function wise fields by molecule, e.g. “coeffs”, “ground_state_coeffs”, “gradient_label”.

__init__(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, batch_size: int = 1, shuffle: bool = False, follow_batch: List[str] | None = ('coeffs', 'atomic_numbers'), exclude_keys: List[str] | None = None, list_keys: List[str] | None = None, **kwargs)

Data loader for OF-DFT data.