diff --git a/sagemaker-train/tests/integ/train/test_dry_run_integration.py b/sagemaker-train/tests/integ/train/test_dry_run_integration.py
index 5133ca7f0a..dcd125a4cb 100644
--- a/sagemaker-train/tests/integ/train/test_dry_run_integration.py
+++ b/sagemaker-train/tests/integ/train/test_dry_run_integration.py
@@ -197,6 +197,7 @@ def test_rlvr_dry_run_returns_none(self, sagemaker_session, valid_dataset):
             training_dataset=valid_dataset,
             accept_eula=True,
         )
+        trainer.hyperparameters.preset_reward_function = "prime_code"
 
         result = trainer.train(dry_run=True)
         assert result is None
@@ -280,6 +281,7 @@ def test_rlvr_serverful_dry_run_returns_none(self, sagemaker_session, valid_data
             compute=TrainingJobCompute(instance_type="ml.g5.12xlarge", instance_count=1),
             accept_eula=True,
         )
+        trainer.hyperparameters.preset_reward_function = "prime_code"
 
         result = trainer.train(dry_run=True)
         assert result is None