following loosely to CUDA doc
@@ -56,15 +56,18 @@ def _get_nvrtc_version():
def _get_arch():
# See Supported Compile Options section of NVRTC User Guide for
# the maximum value allowed for `--gpu-architecture`.
+ # Note that the second-to-max values are used, as the max usually
+ # corresponds to Tegra/Jetson devices
+ # TODO(leofang): figure out a way to support Tegra/Jetson?
major, minor = _get_nvrtc_version()
if major < 9:
# CUDA 8.0
- _nvrtc_max_compute_capability = '62'
- elif major < 10:
- # CUDA 9.x
- _nvrtc_max_compute_capability = '72'
+ _nvrtc_max_compute_capability = '52'
+ elif major < 10 or (major == 10 and minor == 0):
+ # CUDA 9.x / 10.0
+ _nvrtc_max_compute_capability = '70'
else:
- # CUDA 10.x
+ # CUDA 10.1 / 10.2
_nvrtc_max_compute_capability = '75'
return min(device.Device().compute_capability,
