From 23eb65de64e62141258acc258c930c6a63daaf42 Mon Sep 17 00:00:00 2001 From: junyanz Date: Tue, 11 Jul 2017 21:38:14 -0700 Subject: fix the bug when set gpu_ids=1 --- options/base_options.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) (limited to 'options') diff --git a/options/base_options.py b/options/base_options.py index b5b92fb..275c8fc 100644 --- a/options/base_options.py +++ b/options/base_options.py @@ -1,7 +1,7 @@ import argparse import os from util import util - +import torch class BaseOptions(): def __init__(self): @@ -54,6 +54,10 @@ class BaseOptions(): id = int(str_id) if id >= 0: self.opt.gpu_ids.append(id) + + # set gpu ids + if len(self.opt.gpu_ids) > 0: + torch.cuda.set_device(self.opt.gpu_ids[0]) args = vars(self.opt) -- cgit v1.2.3-70-g09d2