Minor fix unit tests

Author: vishwakaria

src/sagemaker/pytorch/estimator.py

@@ -241,6 +241,7 @@ def _pytorch_distribution_configuration(self, distribution):
             dict containing Pytorch DDP config
         """
         distribution_config = {}
+        pytorch_ddp_enabled = False
         if "pytorchddp" in distribution:
             pytorch_ddp_enabled = distribution.get("pytorchddp").get("enabled", False)
 

tests/integ/test_pytorchddp.py

@@ -33,8 +33,8 @@
 )
 def test_pytorchddp_pt_mnist(
     sagemaker_session,
-    pytorch_training_latest_version,
-    pytorch_training_latest_py_version,
+    pytorch_ddp_framework_version,
+    pytorch_ddp_py_version,
 ):
     job_name = sagemaker.utils.unique_name_from_base("pt-pytorch-ddp")
     estimator = PyTorch(
@@ -44,8 +44,8 @@ def test_pytorchddp_pt_mnist(
         instance_count=2,
         instance_type="ml.p3.16xlarge",
         sagemaker_session=sagemaker_session,
-        framework_version=pytorch_training_latest_version,
-        py_version=pytorch_training_latest_py_version,
+        framework_version=pytorch_ddp_framework_version,
+        py_version=pytorch_ddp_py_version,
         distribution={"pytorchddp": {"enabled": True}},
     )
 

tests/unit/test_pytorch.py

@@ -765,6 +765,8 @@ def test_register_pytorch_model_auto_infer_framework(
     sagemaker_session.create_model_package_from_containers.assert_called_with(
         **expected_create_model_package_request
     )
+
+
 def test_pytorch_ddp_distribution_configuration(
     sagemaker_session, pytorch_ddp_framework_version, pytorch_ddp_py_version
 ):