Bring back running pytorch tests in rocm CI

Summary: Pull Request resolved: pytorch#16829

Differential Revision: D13982323

Pulled By: bddppq

fbshipit-source-id: 6ffadb96b9e2ebd64a29e38674a51401dfb211db

Author: bddppq

.jenkins/pytorch/test.sh

@@ -188,27 +188,22 @@ test_xla() {
   assert_git_not_dirty
 }
 
-if [ -z "${BUILD_ENVIRONMENT}" ] || [[ "${BUILD_ENVIRONMENT}" == *-test ]]; then
-  if [[ "${BUILD_ENVIRONMENT}" == *xla* ]]; then
-    test_torchvision
-    test_xla
-  else
-    test_torchvision
-    test_python_nn
-    test_python_all_except_nn
-    test_aten
-    test_libtorch
-    test_custom_script_ops
-  fi
+if [[ "${BUILD_ENVIRONMENT}" == *xla* ]]; then
+  test_torchvision
+  test_xla
+elif [[ "${BUILD_ENVIRONMENT}" == *-test1 ]]; then
+  test_torchvision
+  test_python_nn
+elif [[ "${BUILD_ENVIRONMENT}" == *-test2 ]]; then
+  test_python_all_except_nn
+  test_aten
+  test_libtorch
+  test_custom_script_ops
 else
-  if [[ "${BUILD_ENVIRONMENT}" == *-test1 ]]; then
-    test_torchvision
-    test_python_nn
-  elif [[ "${BUILD_ENVIRONMENT}" == *-test2 ]]; then
-    test_torchvision
-    test_python_all_except_nn
-    test_aten
-    test_libtorch
-    test_custom_script_ops
-  fi
+  test_torchvision
+  test_python_nn
+  test_python_all_except_nn
+  test_aten
+  test_libtorch
+  test_custom_script_ops
 fi