Fixing unit tests for MWMS

Lokiiiiii · Lokiiiiii · commit 5f4959e301ce · 2023-02-10T18:08:51.000Z
diff --git a/tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py b/tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py
@@ -213,7 +213,7 @@ def test_mwms(self, tensorflow_training_version, tensorflow_training_py_version)
                 framework_version=tensorflow_training_version,
                 enable_sagemaker_metrics=False,
                 compiler_config=TrainingCompilerConfig(),
-                distribution={"multi_worker_mirrored_strategy": True},
+                distribution={"multi_worker_mirrored_strategy": {"enabled": True}},
             ).fit()
 
     def test_python_2(self, tensorflow_training_version):
