@@ -37,9 +37,9 @@ def prepare_cuda_env(cuda_version: str, dryrun=False):
37
37
env ["CUDA_HOME" ] = cuda_path_str
38
38
env ["PATH" ] = f"{ cuda_path_str } /bin:{ env ['PATH' ]} "
39
39
env ["CMAKE_CUDA_COMPILER" ] = str (cuda_path .joinpath ("bin" , "nvcc" ).resolve ())
40
- env ["LD_LIBRARY_PATH" ] = (
41
- f" { cuda_path_str } /lib64: { cuda_path_str } /extras/CUPTI/lib64: { env [ ' LD_LIBRARY_PATH' ] } "
42
- )
40
+ env [
41
+ " LD_LIBRARY_PATH"
42
+ ] = f" { cuda_path_str } /lib64: { cuda_path_str } /extras/CUPTI/lib64: { env [ 'LD_LIBRARY_PATH' ] } "
43
43
if dryrun :
44
44
print (f"CUDA_HOME is set to { env ['CUDA_HOME' ]} " )
45
45
# step 2: test call to nvcc to confirm the version is correct
@@ -88,6 +88,7 @@ def setup_cuda_softlink(cuda_version: str):
88
88
89
89
def install_pytorch_nightly (cuda_version : str , env , dryrun = False ):
90
90
from .torch_utils import TORCH_NIGHTLY_PACKAGES
91
+
91
92
uninstall_torch_cmd = ["pip" , "uninstall" , "-y" ]
92
93
uninstall_torch_cmd .extend (TORCH_NIGHTLY_PACKAGES )
93
94
if dryrun :
@@ -168,11 +169,15 @@ def install_torch_deps(cuda_version: str):
168
169
install_torch_deps (cuda_version = args .cudaver )
169
170
if args .install_torch_build_deps :
170
171
from .torch_utils import install_torch_build_deps
172
+
171
173
install_torch_deps (cuda_version = args .cudaver )
172
174
install_torch_build_deps ()
173
175
if args .install_torch_nightly :
174
176
install_pytorch_nightly (cuda_version = args .cudaver , env = os .environ )
175
177
if args .check_torch_nightly_version :
176
178
from .torch_utils import check_torch_nightly_version
177
- assert not args .install_torch_nightly , "Error: Can't run install torch nightly and check version in the same command."
179
+
180
+ assert (
181
+ not args .install_torch_nightly
182
+ ), "Error: Can't run install torch nightly and check version in the same command."
178
183
check_torch_nightly_version (args .force_date )
0 commit comments