def is_available(): global _available if _available is None: _available = False try: _available = runtime.getDeviceCount() > 0 except Exception as e: if (e.args[0] != 'cudaErrorNoDevice: no CUDA-capable device is detected'): raise return _available
from cupy.cuda import cufft from cupy.cuda import device from cupy.cuda import runtime from cupy.fft import config from .test_fft import (multi_gpu_config, _skip_multi_gpu_bug) def intercept_stdout(func): with io.StringIO() as buf, contextlib.redirect_stdout(buf): func() stdout = buf.getvalue() return stdout n_devices = runtime.getDeviceCount() class TestPlanCache(unittest.TestCase): def setUp(self): self.caches = [] self.old_sizes = [] for i in range(n_devices): with device.Device(i): cache = config.get_plan_cache() self.old_sizes.append(cache.get_size()) cache.clear() cache.set_memsize(-1) cache.set_size(2) self.caches.append(cache)