fix installation process

Author: Steboss

.github/container/Dockerfile.axlearn

@@ -25,12 +25,6 @@ seqio==0.0.18
 protobuf==3.20.3  
 pytest>=7.4.3
 REQUIREMENTS
-  # Only append "tensorflow-cpu" if running on x86_64
-  if [ "$(uname -m)" = "x86_64" ]; then
-    echo "tensorflow-cpu" >> /opt/pip-tools.d/requirements-axlearn.in
-  else
-    echo "Skipping TF on $(uname -m)"
-  fi
 EOF
 
 

.github/container/pip-finalize.sh

@@ -46,15 +46,12 @@ if [[ $(echo -n "$unpinned_vcs_dependencies" | wc -l) -gt 0 ]]; then
   exit 1
 fi
 
-# Replace any tensorflow==X with tensorflow-cpu==X in requirements.txt
-sed -i 's/^tensorflow==\([0-9.*]\+\)$/tensorflow-cpu==\1/' requirements.txt
-
-# Replace any torch==Y with torch==Y+cpu in requirements.txt
-sed -i 's/^torch==\([0-9.*]\+\)$/torch==\1+cpu/' requirements.txt
-
-# Add the --find-links option for PyTorch wheels
-echo "--find-links https://download.pytorch.org/whl/torch" >> requirements.txt
-
+# Replace any tensorflow==X with tensorflow-cpu==X in requirements.txt only on amd64
+if [ "$(uname -m)" = "x86_64" ]; then
+  sed -i 's/^tensorflow==\([0-9.*]\+\)$/tensorflow-cpu==\1/' requirements.txt
+else
+  echo "Skipping TF on $(uname -m)"
+fi
 # --no-deps is required since conflicts can still appear during pip-sync
 pip-sync --pip-args '--no-deps --src /opt' requirements.txt
 

.github/container/test-axlearn.sh

@@ -93,8 +93,8 @@ cd "$DIR" || exit 1
 
 echo "Running tests..."
 
-pip install transformers --no-deps 
-pip install scikit-learn timm 
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+pip install timm transformers scikit-learn 
 
 
 if [ "${#TEST_FILES[@]}" -eq 0 ]; then