diff --git a/{{ cookiecutter.repository_name }}/src/{{ cookiecutter.package_name }}/data/dataset.py b/{{ cookiecutter.repository_name }}/src/{{ cookiecutter.package_name }}/data/dataset.py
index 3c63b33..a1ccdd7 100644
--- a/{{ cookiecutter.repository_name }}/src/{{ cookiecutter.package_name }}/data/dataset.py	
+++ b/{{ cookiecutter.repository_name }}/src/{{ cookiecutter.package_name }}/data/dataset.py	
@@ -43,7 +43,7 @@ def main(cfg: omegaconf.DictConfig) -> None:
     Args:
         cfg: the hydra configuration
     """
-    _: MyDataset = hydra.utils.instantiate(cfg.nn.data.datasets.train, split="train", _recursive_=False)
+    _: Dataset = hydra.utils.instantiate(cfg.nn.data.datasets.train, split="train", _recursive_=False)
 
 
 if __name__ == "__main__":