def _init_device(self):
if not torch.cuda.is_available():
self.logger.info('no gpu device available')
sys.exit(1)
np.random.seed(self.cfg.get('seed', 1337))
torch.manual_seed(self.cfg.get('seed', 1337))
torch.cuda.manual_seed(self.cfg.get('seed', 1337))
cudnn.enabled = True
cudnn.benchmark = True
self.device_id, self.gpus_info = get_gpus_memory_info()
self.device = torch.device('cuda:{}'.format(0 if self.cfg['training']['multi_gpus'] else self.device_id))
def _init_device(self):
self.device = torch.device("cuda" if self.cfg['searching']['gpu'] else "cpu")
np.random.seed(self.cfg.get('seed', 1337))
torch.manual_seed(self.cfg.get('seed', 1337))
if self.cfg['searching']['gpu'] and torch.cuda.is_available() :
self.device_id, _ = get_gpus_memory_info()
self.device = torch.device('cuda:{}'.format(0 if self.cfg['searching']['multi_gpus'] else self.device_id))
torch.cuda.manual_seed(self.cfg.get('seed', 1337))
torch.cuda.set_device(self.device)
cudnn.enabled = True
cudnn.benchmark = True
else:
self.logger.info('No gpu devices available!, we will use cpu')
self.device = torch.device('cpu')
self.device_id = 0