fixed logger sample.

shibing624 · Jun 20, 2023 · 9e7c163 · 9e7c163
1 parent 2232bdf
commit 9e7c163
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/textgen/chatglm/chatglm_model.py b/textgen/chatglm/chatglm_model.py
@@ -415,7 +415,7 @@ def train_model(
 
         # Training
         logger.info("*** Train ***")
-        logger.debug(f"Train dataloader example: {list(trainer.get_train_dataloader())[0]}")
+        logger.debug(f"Train dataloader example: {next(iter(trainer.get_train_dataloader()))}")
         (global_step, training_loss, metrics) = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
         self.results.update(metrics)
         trainer.log_metrics("train", metrics)

diff --git a/textgen/gpt/gpt_model.py b/textgen/gpt/gpt_model.py
@@ -446,7 +446,7 @@ def train_model(
 
         # Training
         logger.info("*** Train ***")
-        logger.debug(f"Train dataloader example: {list(trainer.get_train_dataloader())[0]}")
+        logger.debug(f"Train dataloader example: {next(iter(trainer.get_train_dataloader()))}")
         (global_step, training_loss, metrics) = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
         self.results.update(metrics)
         trainer.log_metrics("train", metrics)