Skip to content

Commit 7c46f8a

Browse files
author
Rohan Gujarathi
committed
remove hardcoded pytorch framework and py_version
1 parent d596f11 commit 7c46f8a

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

tests/integ/test_workflow.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,8 @@ def test_training_job_with_debugger(
413413
sagemaker_session,
414414
pipeline_name,
415415
role,
416+
pytorch_training_latest_version,
417+
pytorch_training_latest_py_version,
416418
):
417419
instance_count = ParameterInteger(name="InstanceCount", default_value=1)
418420
instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
@@ -424,6 +426,7 @@ def test_training_job_with_debugger(
424426
]
425427
debugger_hook_config = DebuggerHookConfig(
426428
s3_output_path=f"s3://{sagemaker_session.default_bucket()}/{uuid.uuid4()}/tensors"
429+
)
427430

428431
base_dir = os.path.join(DATA_DIR, "pytorch_mnist")
429432
script_path = os.path.join(base_dir, "mnist.py")
@@ -436,8 +439,8 @@ def test_training_job_with_debugger(
436439
pytorch_estimator = PyTorch(
437440
entry_point=script_path,
438441
role="SageMakerRole",
439-
framework_version="1.5.0",
440-
py_version="py3",
442+
framework_version=pytorch_training_latest_version,
443+
py_version=pytorch_training_latest_py_version,
441444
instance_count=instance_count,
442445
instance_type=instance_type,
443446
sagemaker_session=sagemaker_session,
@@ -462,7 +465,7 @@ def test_training_job_with_debugger(
462465
response = pipeline.create(role)
463466
create_arn = response["PipelineArn"]
464467

465-
execution = pipeline.start(parameters={})
468+
execution = pipeline.start()
466469
response = execution.describe()
467470
assert response["PipelineArn"] == create_arn
468471

0 commit comments

Comments
 (0)