File size: 548 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | from typing import Dict
from omegaconf import DictConfig
from datasets.video.base_video import BaseAdvancedVideoDataset
from . import BaseDataModule, ResumableDataModule
def _data_module_cls(_, root_cfg: DictConfig, compatible_datasets: Dict):
dataset_cls = compatible_datasets[root_cfg.dataset._name]
data_module_cls = None
if issubclass(dataset_cls, BaseAdvancedVideoDataset):
data_module_cls = ResumableDataModule
else:
data_module_cls = BaseDataModule
return data_module_cls(root_cfg, compatible_datasets)
|