Skip to content

Commit ce4a2e3

Browse files
committed
first try
1 parent 770311c commit ce4a2e3

File tree

1 file changed

+8
-0
lines changed

1 file changed

+8
-0
lines changed

test/smoke_test/smoke_test.py

+8
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,14 @@ def smoke_test_cuda() -> None:
1919
print(f"torch cuda: {torch.version.cuda}")
2020
# todo add cudnn version validation
2121
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+
2230

2331
def smoke_test_conv2d() -> None:
2432
import torch.nn as nn

0 commit comments

Comments
 (0)