Class DataModuleTemplateFromConfig

main.py:192–203

class DataModuleTemplateFromConfig(DataModuleFromConfig):
    def __init__(self, batch_size, train=None, validation=None, test=None,
                 wrap=False, num_workers=None):
        super().__init__(batch_size, train=train, validation=validation,
                         test=test, wrap=wrap, num_workers=num_workers)
        if train is not None:
            self.dataset_configs["train"] = train
            self.train_dataloader = self._train_dataloader

    def _train_dataloader(self):
        return DataLoader(self.datasets["train"], batch_size=self.batch_size*2,
                        num_workers=self.num_workers, shuffle=True, collate_fn=custom_collate)


class SetupCallback(Callback):