Updated test case based on recent commits

Author: quic-meetkuma

tests/finetune/test_finetune.py

@@ -66,21 +66,21 @@ def test_finetune(
 
     results = finetune(**kwargs)
 
-    assert np.allclose(results["avg_train_prep"], 1.002326), "Train perplexity is not matching."
-    assert np.allclose(results["avg_train_loss"], 0.00232327), "Train loss is not matching."
-    assert np.allclose(results["avg_eval_prep"], 1.0193923), "Eval perplexity is not matching."
-    assert np.allclose(results["avg_eval_loss"], 0.0192067), "Eval loss is not matching."
-    assert results["avg_epoch_time"] < 30, "Training should complete within 30 seconds."
+    assert np.allclose(results["avg_train_prep"], 1.002326, atol=1e-5), "Train perplexity is not matching."
+    assert np.allclose(results["avg_train_loss"], 0.00232327, atol=1e-5), "Train loss is not matching."
+    assert np.allclose(results["avg_eval_prep"], 1.0193923, atol=1e-5), "Eval perplexity is not matching."
+    assert np.allclose(results["avg_eval_loss"], 0.0192067, atol=1e-5), "Eval loss is not matching."
+    assert results["avg_epoch_time"] < 60, "Training should complete within 60 seconds."
 
     train_config_spy.assert_called_once()
     generate_dataset_config_spy.assert_called_once()
     generate_peft_config_spy.assert_called_once()
-    update_config_spy.assert_called_once()
     get_custom_data_collator_spy.assert_called_once()
     get_longest_seq_length_spy.assert_called_once()
     print_model_size_spy.assert_called_once()
     train_spy.assert_called_once()
 
+    assert update_config_spy.call_count == 2
     assert get_dataloader_kwargs_spy.call_count == 2
     assert get_preprocessed_dataset_spy.call_count == 2
 
@@ -102,7 +102,7 @@ def test_finetune(
     else:
         assert eval_dataloader is None
 
-    args, kwargs = update_config_spy.call_args
+    args, kwargs = update_config_spy.call_args_list[0]
     train_config = args[0]
     assert max_train_step >= train_config.gradient_accumulation_steps, (
         "Total training step should be more than 4 which is gradient accumulation steps."