diff --git a/src/tf_container/train_entry_point.py b/src/tf_container/train_entry_point.py
index 1541019d..bc0365ec 100644
--- a/src/tf_container/train_entry_point.py
+++ b/src/tf_container/train_entry_point.py
@@ -137,6 +137,7 @@ def train():
     os.environ['S3_REQUEST_TIMEOUT_MSEC'] = str(env.hyperparameters.get('s3_checkpoint_save_timeout', 60000))
 
     env.download_user_module()
+    env.pip_install_requirements()
 
     customer_script = env.import_user_module()
 
diff --git a/test/integ/test_estimator_classification.py b/test/integ/test_estimator_classification.py
index c66645e3..479f5307 100644
--- a/test/integ/test_estimator_classification.py
+++ b/test/integ/test_estimator_classification.py
@@ -23,7 +23,6 @@
 def test_estimator_classification(docker_image, sagemaker_session, opt_ml, processor):
     resource_path = os.path.join(SCRIPT_PATH, '../resources/iris')
 
-    copy_resource(resource_path, opt_ml, 'code')
     copy_resource(resource_path, opt_ml, 'data', 'input/data')
 
     s3_source_archive = fw_utils.tar_and_upload_dir(session=sagemaker_session.boto_session,