fix(tests): fix ci tests error

InternLM · Jan 29, 2024 · f02523e · f02523e
1 parent 011edcf
commit f02523e
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/internlm/train/training_internlm.py b/internlm/train/training_internlm.py
@@ -78,11 +78,11 @@
 from internlm.utils.logger import get_logger
 from internlm.utils.megatron_timers import megatron_timer as timer
 from internlm.utils.parallel import (
-    is_using_isp,
     is_replica_zero_parallel_parameter,
     is_tensor_data_parallel_parameter,
     is_tensor_expert_data_parallel_parameter,
     is_tensor_zero_parallel_parameter,
+    is_using_isp,
     is_weight_zero_parallel_parameter,
     set_model_params_layer_name,
     sync_model_param,

diff --git a/tests/test_training/test_loss.py b/tests/test_training/test_loss.py
@@ -93,7 +93,7 @@ def train(
     current_time = objs[0]
 
     # initialize model
-    model, _ = initialize_model()
+    model = initialize_model()
 
     # initialize loss function
     criterion = FlashGPTLMLoss(parallel_output=True, label_smoothing=label_smoothing)

diff --git a/tests/test_training/test_swap_nb_loss_and_gradnorm.py b/tests/test_training/test_swap_nb_loss_and_gradnorm.py
@@ -278,7 +278,7 @@ def exam_loss(args):
     seed_all(1024)
 
     # initialize model
-    model, _ = initialize_model()
+    model = initialize_model()
 
     # initialize loss function
     criterion = FlashGPTLMLoss(parallel_output=True, label_smoothing=gpc.config.loss.label_smoothing)

diff --git a/tests/test_training/train_CI.py b/tests/test_training/train_CI.py
@@ -124,7 +124,7 @@ def main(args):
     uniscale_logger = initialize_llm_logger(start_time=current_time)
 
     # initialize model
-    model, _ = initialize_model()
+    model = initialize_model()
 
     with open(args.config, "r") as f:
         config_lines = f.readlines()