Update build-extension to allow arm compilation

nemotron-ocr/scripts/build-extension.py

@@ -5,6 +5,7 @@
 
 
 import os
+import platform
 import sys
 import shutil
 from glob import glob
@@ -73,8 +74,8 @@ common_args = ["-std=c++17"]
 gcc_args = ["-fopenmp", "-DOMP_NESTED=true"]
 cuda_args = []
 
-# TODO: analyze if SIMD optimizations are beneficial on ARM
-if os.environ.get("ARCH") != "arm64":
+_target_arch = os.environ.get("ARCH", "").lower() or platform.machine().lower()
+if _target_arch in ("x86_64", "amd64"):
     gcc_args += ["-mavx2"]
 
 files = sorted(glob("cpp/**/*.cu", recursive=True) + glob("cpp/**/*.cpp", recursive=True))