from ._registry import DATALOADERS

from torch.utils.data import DataLoader


@DATALOADERS.register
class DataListLoader(DataLoader):
    r"""Data loader which merges data objects from a
    :class:`torch_geometric.data.dataset` to a python list.
    .. note::
        This data loader should be used for multi-gpu support via
        :class:`torch_geometric.nn.DataParallel`.
    Args:
        dataset (Dataset): The dataset from which to load the data.
        batch_size (int, optional): How many samples per batch to load.
            (default: :obj:`1`)
        shuffle (bool, optional): If set to :obj:`True`, the data will be
            reshuffled at every epoch (default: :obj:`False`)
    """

    def __init__(self, dataset, batch_size=1, shuffle=False, **kwargs):
        super().__init__(
            dataset, batch_size, shuffle,
            collate_fn=lambda data_list: data_list, **kwargs)

from ._registry import DATALOADERS

from torch.utils.data import DataLoader


@DATALOADERS.register
class DataListLoader(DataLoader):
    r"""Data loader which merges data objects from a
    :class:`torch_geometric.data.dataset` to a python list.
    .. note::
        This data loader should be used for multi-gpu support via
        :class:`torch_geometric.nn.DataParallel`.
    Args:
        dataset (Dataset): The dataset from which to load the data.
        batch_size (int, optional): How many samples per batch to load.
            (default: :obj:`1`)
        shuffle (bool, optional): If set to :obj:`True`, the data will be
            reshuffled at every epoch (default: :obj:`False`)
    """

    def __init__(self, dataset, batch_size=1, shuffle=False, **kwargs):
        super().__init__(
            dataset, batch_size, shuffle,
            collate_fn=lambda data_list: data_list, **kwargs)

Data loader which merges data objects from a     :class:`torch_geometric.data.dataset` to a python list.     .. note::         This data loader should be used for multi-gpu support via         :class:`torch_geometric.nn.DataParallel`.     Args:         dataset (Dataset): The dataset from which to load the data.         batch_size (int, optional): How many samples per batch to load.             (default: :obj:`1`)         shuffle (bool, optional): If set to :obj:`True`, the data will be             reshuffled at every epoch (default: :obj:`False`)