Commit ce08f8d2 authored by hwangjeff's avatar hwangjeff Committed by Facebook GitHub Bot
Browse files

Fix CUDA check (#2710)

Summary:
`torch.version.cuda` can return a string of form X.X or X.X.X. This PR modifies the CUDA version check to account for this.

Pull Request resolved: https://github.com/pytorch/audio/pull/2710

Reviewed By: carolineechen, nateanl

Differential Revision: D39796810

Pulled By: hwangjeff

fbshipit-source-id: b483bd8200195844d65d0caddebaf1b10f939b64
parent 0a5825ae
......@@ -105,7 +105,8 @@ def _check_cuda_version():
if version is not None and torch.version.cuda is not None:
version_str = str(version)
ta_version = f"{version_str[:-3]}.{version_str[-2]}"
t_version = torch.version.cuda
t_version = torch.version.cuda.split(".")
t_version = f"{t_version[0]}.{t_version[1]}"
if ta_version != t_version:
raise RuntimeError(
"Detected that PyTorch and TorchAudio were compiled with different CUDA versions. "
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment