mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
9 lines
209 B
Python
9 lines
209 B
Python
import os
|
|
|
|
from .env import check_env_flag
|
|
|
|
CUDA_HOME = os.getenv('CUDA_HOME', '/usr/local/cuda')
|
|
WITH_CUDA = not check_env_flag('NO_CUDA') and os.path.exists(CUDA_HOME)
|
|
if not WITH_CUDA:
|
|
CUDA_HOME = None
|