lol

vall_e/engines/base.py

@@ -446,7 +446,7 @@ class Engines(dict[str, Engine]):
 				engine.tokens_processed = 0
 
 		# update the LR because for some god awful reason it gets overwritten when loading from a checkpoint but only when it's not using a scheduler
-		if cfg.hyperparameters.scheduler_type == "":
+		if cfg.hyperparameters.scheduler == "":
 			self.set_lr(cfg.hyperparameters.learning_rate)
 
 		self._update()

vall_e/engines/deepspeed.py

@@ -47,7 +47,7 @@ class Engine(DeepSpeedEngine):
 		}
 
 		# kwargs['stats'] = None will return None when popped
-		maybe_stats = kwargs.get('stats', stats)
+		maybe_stats = kwargs.pop('stats', stats)
 		if maybe_stats is not None:
 			stats = maybe_stats