Update fw_utils.py - support 2.0.1 container for DDP and Torch distri… (#4127)

Author: gilinachum

src/sagemaker/fw_utils.py

@@ -151,10 +151,11 @@
     "1.12.1",
     "1.13.1",
     "2.0.0",
+    "2.0.1",
 ]
 
 
-TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = ["1.13.1", "2.0.0"]
+TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = ["1.13.1", "2.0.0", "2.0.1"]
 
 TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
 TRAINIUM_SUPPORTED_TORCH_DISTRIBUTED_FRAMEWORK_VERSIONS = [