reticulating splines

vall_e/engines/__init__.py

@@ -5,5 +5,4 @@ if cfg.trainer.backend == "deepspeed":
 elif cfg.trainer.backend == "local":
 	from .base import Engine
 
-from .base import Engines
-from .base import TrainFeeder
+from .base import Engines, TrainFeeder, default_feeder

vall_e/engines/base.py

@@ -9,7 +9,7 @@ class TrainFeeder(Protocol):
 	) -> None | tuple[Tensor, Stats]:
 		...
 
-def default_feed(engine, batch):
+def default_feeder(engine, batch):
 	if isinstance(batch, list):
 		engine( *batch )
 	elif isinstance(batch, dict):
@@ -257,7 +257,7 @@ class Engines(dict[str, Engine]):
 			stats.update(flatten_dict({ name.split("-")[0]: stat }))
 		return stats
 
-	def step(self, batch, feeder: TrainFeeder = default_feed, device=torch.cuda.current_device()):
+	def step(self, batch, feeder: TrainFeeder = default_feeder, device=torch.cuda.current_device()):
 		total_elapsed_time = 0
 
 		stats: Any = dict()

vall_e/models/__init__.py

@@ -11,10 +11,10 @@ def get_model(cfg):
 	name = cfg.name
 
 	model = Model(
-		n_tokens=model.tokens,
-		d_model=model.dim,
-		n_heads=model.heads,
-		n_layers=model.layers,
+		n_tokens=cfg.tokens,
+		d_model=cfg.dim,
+		n_heads=cfg.heads,
+		n_layers=cfg.layers,
 	)
 	model._cfg = cfg