@@ -196,10 +196,7 @@ def build_train_config(config_path: Path, run_name: str, group_id: str, beaker_u
196196trainer = config .trainer .build (train_module ,data_loader )
197197
198198# If we have a load path and there is no checkpoint in the save folder, load the checkpoint from the load path.
199- if (
200- not trainer .maybe_load_checkpoint (trainer .save_folder ,load_trainer_state = base_config .load_state )
201- and base_config .load_path
202- ):
199+ if not trainer .maybe_load_checkpoint (trainer .save_folder )and base_config .load_path :
203200logger .info (
204201f"Loading checkpoint from{ base_config .load_path } { base_config .load_state }
205202 )
@@ -241,7 +238,7 @@ def mk_launch_configs(group: ExperimentGroup, beaker_user: str) -> list[BeakerLa
241238budget = group .config .budget or "ai2/oe-data" ,
242239workspace = group .config .workspace ,
243240preemptible = group .config .preemptible ,
244- beaker_image = "petew/olmo-core-tch270cu126 " ,
241+ beaker_image = "petew/olmo-core-tch270cu128 " ,
245242priority = group .config .priority ,
246243env_vars = [BeakerEnvVar (name = "NCCL_DEBUG" ,value = "INFO" if group .config .nccl_debug else "WARN" )],
247244env_secrets = [