Fix post-script on Nova

Differential Revision: D52377515

.github/scripts/nova_postscript.bash

@@ -28,13 +28,22 @@ install_fbgemm_gpu_wheel "${BUILD_ENV_NAME}" fbgemm_gpu/dist/*.whl
 
 # Test with PyTest
 echo "[NOVA] Current working directory: $(pwd)"
-CPU_GPU="${CU_VERSION}"
-if [ "${CU_VERSION}" != 'cpu' ]; then
-    CPU_GPU=""
+if [[ $CU_VERSION = cu* ]]; then
+  echo "[NOVA] Testing the CUDA variant of FBGEMM_GPU ..."
+  export fbgemm_variant="cuda"
+
+elif [[ $CU_VERSION = rocm* ]]; then
+  echo "[NOVA] Testing the ROCm variant of FBGEMM_GPU ..."
+  export fbgemm_variant="rocm"
+
+else
+  echo "[NOVA] Testing the CPU variant of FBGEMM_GPU ..."
+  export fbgemm_variant="cpu"
 fi
+
 $CONDA_RUN python3 -c "import torch; print('cuda.is_available() ', torch.cuda.is_available()); print ('device_count() ',torch.cuda.device_count());"
 cd "${FBGEMM_REPO}/fbgemm_gpu/test" || { echo "[NOVA] Failed to cd to fbgemm_gpu/test from $(pwd)"; };
-run_fbgemm_gpu_tests "${BUILD_ENV_NAME}" "${CPU_GPU}"
+run_fbgemm_gpu_tests "${BUILD_ENV_NAME}" "${fbgemm_variant}"
 
 # Workaround EACCES: permission denied error at checkout step
 chown -R 1000:1000 /__w/FBGEMM/FBGEMM/ || echo "Unable to chown 1000:1000 from $USER, uid: $(id -u)"