fixed the integration test

Author: ChaoPang

src/cehrbert/runners/hf_cehrbert_finetune_runner.py

@@ -27,7 +27,6 @@
 )
 from cehrbert.data_generators.hf_data_generator.hf_dataset_mapping import MedToCehrBertDatasetMapping
 from cehrbert.data_generators.hf_data_generator.meds_utils import create_dataset_from_meds_reader
-from cehrbert.data_generators.hf_data_generator.sample_packing_sampler import SamplePackingBatchSampler
 from cehrbert.models.hf_models.config import CehrBertConfig
 from cehrbert.models.hf_models.hf_cehrbert import CehrBertForClassification, CehrBertLstmForClassification
 from cehrbert.models.hf_models.tokenization_hf_cehrbert import CehrBertTokenizer

tests/integration_tests/runners/hf_cehrbert_pretrain_runner_test.py

@@ -49,6 +49,16 @@ def test_train_model(self):
             self.dataset_prepared_path,
             "--max_steps",
             "10",
+            "--save_strategy",
+            "steps",
+            "--eval_strategy",
+            "steps",
+            "--do_train",
+            "true",
+            "--do_predict",
+            "true",
+            "--load_best_model_at_end",
+            "true",
             "--report_to",
             "none",
         ]