find cuda home

app.py

@@ -1,6 +1,23 @@
 import os
 import subprocess
 
+# 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'
+]
+
+cuda_home = None
+for path in cuda_paths:
+    if os.path.exists(path):
+        cuda_home = path
+        break
+
+if cuda_home is None:
+    raise EnvironmentError('CUDA installation not found. Please install CUDA or set CUDA_HOME manually.')
+
 # Set the CUDA_HOME environment variable
 os.environ['CUDA_HOME'] = '/usr/local/cuda'
 os.environ['PATH'] = os.environ['CUDA_HOME'] + '/bin:' + os.environ['PATH']