fix a runaway weights folder

hkproj · Sep 16, 2023 · c8f44fb · c8f44fb
1 parent 3cf3af2
commit c8f44fb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/config.py b/config.py
@@ -25,8 +25,7 @@ def get_weights_file_path(config, epoch: str):
 # Find the latest weights file in the weights folder
 def latest_weights_file_path(config):
     model_folder = f"{config['datasource']}_{config['model_folder']}"
-    model_basename = config["model_basename"]
-    model_filename = f"{config['model_basename']}*.pt"
+    model_filename = f"{config['model_basename']}*"
     weights_files = list(Path(model_folder).glob(model_filename))
     if len(weights_files) == 0:
         return None

diff --git a/train.py b/train.py
@@ -183,7 +183,7 @@ def train_model(config):
     print("Using device:", device)
 
     # Make sure the weights folder exists
-    Path(config['model_folder']).mkdir(parents=True, exist_ok=True)
+    Path(f"{config['datasource']}_{config['model_folder']}").mkdir(parents=True, exist_ok=True)
 
     train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt = get_ds(config)
     model = get_model(config, tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size()).to(device)