if config["cuda"] and not torch.cuda.is_available():
self.logger.warning("Warning: There\"s no GPU available on this machine, "
"training will be performed on CPU.")
self.device = torch.device("cuda:" + str(config["gpu"]) if self.with_cuda else "cpu")
self.model = model.to(self.device)
self.loss = loss
After Change
n_gpu_use = 0
if n_gpu_use > n_gpu:
msg = "Warning: The number of GPU\"s configured to use is {}, but only {} are available on this machine.".format(n_gpu_use, n_gpu)
self.logger.warning(msg)n_gpu_use = n_gpu
self.device = torch.device("cuda:0" if n_gpu_use > 0 else "cpu")
self.model = model.to(self.device)
if n_gpu_use > 1:
self.model = torch.nn.DataParallel(model, device_ids=list(range(n_gpu_use)))
self.loss = loss
self.metrics = metrics
