code logging

viraweb123 · Aug 14, 2024 · f4bbf06 · f4bbf06
1 parent 80d9fb4
commit f4bbf06
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/train/code/main.py b/train/code/main.py
@@ -12,6 +12,7 @@
 input_base_dir = "/train/input"    
 
 output_dir = os.path.join(output_base_dir, log_dir_name, "trained_model")
+logging_dir = os.path.join(output_base_dir, "model_log")
 log_dir = os.path.join(log_base_dir, log_dir_name)
 
 os.makedirs(output_dir, exist_ok=True)
@@ -81,6 +82,8 @@ def train(train_file_path,
         overwrite_output_dir=overwrite_output_dir,
         per_device_train_batch_size=per_device_train_batch_size,
         num_train_epochs=num_train_epochs,
+        logging_dir=logging_dir, 
+        logging_steps=10,
         save_steps=save_steps
     )
     trainer = Trainer(