Skip to content

Commit 2c5212e

Browse files
committed
Minor fix unit tests
1 parent b91ac25 commit 2c5212e

File tree

2 files changed

+5
-4
lines changed

2 files changed

+5
-4
lines changed

src/sagemaker/pytorch/estimator.py

+1
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,7 @@ def _pytorch_distribution_configuration(self, distribution):
241241
dict containing Pytorch DDP config
242242
"""
243243
distribution_config = {}
244+
pytorch_ddp_enabled = False
244245
if "pytorchddp" in distribution:
245246
pytorch_ddp_enabled = distribution.get("pytorchddp").get("enabled", False)
246247

tests/integ/test_pytorchddp.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@
3333
)
3434
def test_pytorchddp_pt_mnist(
3535
sagemaker_session,
36-
pytorch_training_latest_version,
37-
pytorch_training_latest_py_version,
36+
pytorch_ddp_framework_version,
37+
pytorch_ddp_py_version,
3838
):
3939
job_name = sagemaker.utils.unique_name_from_base("pt-pytorch-ddp")
4040
estimator = PyTorch(
@@ -44,8 +44,8 @@ def test_pytorchddp_pt_mnist(
4444
instance_count=2,
4545
instance_type="ml.p3.16xlarge",
4646
sagemaker_session=sagemaker_session,
47-
framework_version=pytorch_training_latest_version,
48-
py_version=pytorch_training_latest_py_version,
47+
framework_version=pytorch_ddp_framework_version,
48+
py_version=pytorch_ddp_py_version,
4949
distribution={"pytorchddp": {"enabled": True}},
5050
)
5151

0 commit comments

Comments
 (0)