From 8eca1a44dd66613b2c614feec3a4e4175f369b0e Mon Sep 17 00:00:00 2001 From: stevenfontanella Date: Sat, 28 Sep 2019 11:52:14 -0400 Subject: [PATCH] Fix crash when running on CPU --- src/config.py | 2 +- src/run.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/config.py b/src/config.py index 07fe97e..5148911 100644 --- a/src/config.py +++ b/src/config.py @@ -12,7 +12,7 @@ def parse_args(): parser.add_argument('--train', action='store_true', help='train the segmentation model') parser.add_argument('--evaluate', action='store_true', help='evaluate the model') parser.add_argument('--segment', action='store_true', help='segment new files or input text') - parser.add_argument('--gpu', type=str, default='0', help='specify gpu device') + parser.add_argument('--gpu', type=str, help='specify gpu device') train_settings = parser.add_argument_group('train settings') train_settings.add_argument('--optim', default='adam', help='optimizer type') diff --git a/src/run.py b/src/run.py index e07188a..124bede 100644 --- a/src/run.py +++ b/src/run.py @@ -32,7 +32,8 @@ logger.info('Running with args : {}'.format(args)) os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu + if args.gpu is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu random.seed(args.seed) np.random.seed(args.seed)