try to find cuda home

app.py

@@ -1,25 +1,19 @@
 import os
 import subprocess
+import glob
 
-# Check common CUDA installation paths and set the correct one
-cuda_paths = [
-    '/usr/local/cuda',  # Default path
-    '/usr/local/cuda-11.0',  # Example of a versioned path
-    '/usr/local/cuda-11.1',
-    '/usr/local/cuda-11.2'
-]
+# Find all CUDA directories that match /usr/local/cuda*
+cuda_dirs = glob.glob('/usr/local/cuda*')
 
-cuda_home = None
-for path in cuda_paths:
-    if os.path.exists(path):
-        cuda_home = path
-        break
+if not cuda_dirs:
+    raise EnvironmentError('No CUDA installation found. Please install CUDA or set CUDA_HOME manually.')
 
-if cuda_home is None:
-    raise EnvironmentError('CUDA installation not found. Please install CUDA or set CUDA_HOME manually.')
+# Assume the highest version of CUDA is the one to use
+cuda_dirs.sort()
+cuda_home = cuda_dirs[-1]
 
 # Set the CUDA_HOME environment variable
-os.environ['CUDA_HOME'] = '/usr/local/cuda'
+os.environ['CUDA_HOME'] = cuda_home
 os.environ['PATH'] = os.environ['CUDA_HOME'] + '/bin:' + os.environ['PATH']
 os.environ['LD_LIBRARY_PATH'] = os.environ['CUDA_HOME'] + '/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')