def decide_cuda(opt: Optional[BaseOptions] = None):
if opt is None:
opt = BaseOptions()
if opt.use_cpu:
return False
def get_error_str(name, err):
e_str = ("Failed to initialize %s library; "
"falling back to CPU. Set 'use_cpu' to "
"True to avoid this warning." % (name))
if err is not None:
e_str += "\nError encountered was %s" % (err)
return e_str
if not torch.cuda.is_available():
warnings.warn(get_error_str("CUDA", None))
return False
try:
from falkon.cuda import cublas_gpu # noqa F401
except Exception as e:
warnings.warn(get_error_str("cuBLAS", e))
return False
try:
from falkon.cuda import cudart_gpu # noqa F401
except Exception as e:
warnings.warn(get_error_str("cudart", e))
return False
try:
from falkon.cuda import cusolver_gpu # noqa F401
except Exception as e:
warnings.warn(get_error_str("cuSOLVER", e))
return False
return True
def decide_cuda(opt: BaseOptions = BaseOptions()):
if opt.use_cpu:
return False
def get_error_str(name, err):
return ("Failed to initialize %s library; "
"falling back to CPU. Set 'use_cpu' to "
"True to avoid this warning."
"\nError encountered was %s" % (name, err))
try:
from falkon.cuda import cublas_gpu
except Exception as e:
warnings.warn(get_error_str("cuBLAS", e))
return False
try:
from falkon.cuda import cudart_gpu
except Exception as e:
warnings.warn(get_error_str("cudart", e))
return False
try:
from falkon.cuda import cusolver_gpu
except Exception as e:
warnings.warn(get_error_str("cuSOLVER", e))
return False
return True
def initialize_cuda():
# your setup code goes here, executed ahead of first test
opt = BaseOptions(compute_arch_speed=False, use_cpu=False)
if decide_cuda():
initialization.init(opt)
def _setup_opt(opt: Optional[BaseOptions], is_cpu=False) -> BaseOptions:
if opt is None:
opt = BaseOptions()
return dataclasses.replace(opt, use_cpu=is_cpu)