abreza's picture
refactor: Update image generation pipeline to use playground v2.5
4e92ab0
raw
history blame
353 Bytes
import os
import shutil
def find_cuda():
cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
if cuda_home and os.path.exists(cuda_home):
return cuda_home
nvcc_path = shutil.which('nvcc')
if nvcc_path:
cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
return cuda_path
return None