File tree Expand file tree Collapse file tree 1 file changed +8
-0
lines changed Expand file tree Collapse file tree 1 file changed +8
-0
lines changed Original file line number Diff line number Diff line change @@ -19,6 +19,14 @@ def smoke_test_cuda() -> None:
19
19
print (f"torch cuda: { torch .version .cuda } " )
20
20
# todo add cudnn version validation
21
21
print (f"torch cudnn: { torch .backends .cudnn .version ()} " )
22
+ # check torchvision's cuda version against system cuda version
23
+ if (torch .ops .torchvision .__cuda__version () != gpu_arch_ver ):
24
+ raise RuntimeError (f"Wrong CUDA version. Loaded: { torch .version .cuda } Expected: { gpu_arch_ver } " )
25
+ # check torchaudio's cuda version against system cuda version
26
+ if 'cu' + str (gpu_arch_ver ).replace ("." , "" ) not in torchaudio .__version__ .split ("+" )
27
+ raise RuntimeError (f"Wrong CUDA version. Loaded: { torchaudio .__version__ } Expected: { gpu_arch_ver }
28
+
29
+
22
30
23
31
def smoke_test_conv2d () -> None :
24
32
import torch .nn as nn
You can’t perform that action at this time.
0 commit comments