Fix FSDP venv python version

Author: mhuguesaws

3.test_cases/pytorch/FSDP/slurm/create_venv.sh

@@ -4,11 +4,22 @@ set -ex
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: MIT-0
 
-PYTHON_VERSION=$(python3 --version | awk '{print $2}' | awk -F'.' '{print $1"."$2}')
+PYTHON_V=python3
+OS_VERSION=$(cat /etc/os-release | grep VERSION_ID | awk -F '=' '{print $2}')
+OS_VERSION=${OS_VERSION//\"/}
+
+if [ $OS_VERSION = "20.04" ]; then
+
+   PYTHON_VERSION=$(python3.9 --version | awk '{print $2}' | awk -F'.' '{print $1"."$2}')
+   PYTHON_V=python3.9
+else
+  PYTHON_VERSION=$(python3 --version | awk '{print $2}' | awk -F'.' '{print $1"."$2}')
+fi
+
 sudo apt install -y python$PYTHON_VERSION-venv
 
 # Create and actiate Python virtual environment
-python3 -m venv env
+$PYTHON_V -m venv env
 source ./env/bin/activate
 
 pip install -r ../src/requirements.txt