diff --git a/easybuild/easyblocks/p/pytorch.py b/easybuild/easyblocks/p/pytorch.py index 862afbf61de..4bb27ea6003 100755 --- a/easybuild/easyblocks/p/pytorch.py +++ b/easybuild/easyblocks/p/pytorch.py @@ -509,14 +509,20 @@ def add_enable_option(name, enabled): # determine the compute capability of a GPU in the system and use that which may fail tests if # it is to new for the used nvcc env.setvar('TORCH_CUDA_ARCH_LIST', ';'.join(cuda_cc)) + with_gpu_support = True else: # Disable CUDA options.append('USE_CUDA=0') + with_gpu_support = False if pytorch_version >= '2.0': - add_enable_option('USE_ROCM', get_software_root('ROCm')) - elif pytorch_version >= 'v1.10.0': - add_enable_option('USE_MAGMA', get_software_root('magma')) + has_rocm = get_software_root('ROCm') + add_enable_option('USE_ROCM', has_rocm) + if has_rocm: + with_gpu_support = True + + if pytorch_version >= '1.10.0': + add_enable_option('USE_MAGMA', with_gpu_support and get_software_root('magma')) if get_cpu_architecture() == POWER: # *NNPACK is not supported on Power, disable to avoid warnings @@ -525,7 +531,7 @@ def add_enable_option(name, enabled): if pytorch_version >= '1.10.0' and pytorch_version < '1.12.0': options.append('USE_BREAKPAD=0') # FBGEMM requires AVX512, so not available on PPC - if pytorch_version >= 'v1.10.0': + if pytorch_version >= '1.10.0': options.append('USE_FBGEMM=0') # Metal only supported on IOS which likely doesn't work with EB, so disabled