diff --git a/cli/src/pcluster/resources/imagebuilder/parallelcluster_test.yaml b/cli/src/pcluster/resources/imagebuilder/parallelcluster_test.yaml
index 56e620a362..0e51efdaa0 100644
--- a/cli/src/pcluster/resources/imagebuilder/parallelcluster_test.yaml
+++ b/cli/src/pcluster/resources/imagebuilder/parallelcluster_test.yaml
@@ -252,7 +252,7 @@ phases:
               echo "Correctly installed CUDA ${cuda_output}"
 
               echo "Testing CUDA with deviceQuery..."
-              if [ {{ test.OSArchitecture.outputs.stdout }} != 'arm64' ]; then
+              if [ -f /usr/local/cuda-${cuda_ver}/extras/demo_suite/deviceQuery ]; then
                 /usr/local/cuda-${cuda_ver}/extras/demo_suite/deviceQuery | grep -o "Result = PASS"
                 [[ $? -ne 0 ]] && echo "CUDA deviceQuery test failed" && exit 1
               else