diff --git a/src/hyfi/main/main.py b/src/hyfi/main/main.py
index f5b29d19..1b6b9948 100644
--- a/src/hyfi/main/main.py
+++ b/src/hyfi/main/main.py
@@ -10,7 +10,7 @@
 from omegaconf import DictConfig, OmegaConf
 
 from hyfi.cached_path import cached_path
-from hyfi.composer import Composer, SpecialKeys
+from hyfi.composer import GENERATOR, Composer, SpecialKeys
 from hyfi.copier import Copier
 from hyfi.core import (
     __app_version__,
@@ -84,6 +84,7 @@ class HyFI(
     DATASETs,
     ENVs,
     FUNCs,
+    GENERATOR,
     GPUs,
     GRAPHICs,
     IOLIBs,