From 4f2903eb8f61b1d6be90df295b5fd8d666466d44 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Wed, 25 Apr 2018 17:01:05 -0600 Subject: [PATCH] Fix a bug for GPU --- network/network_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/network/network_utils.py b/network/network_utils.py index fca5135..c3659a1 100644 --- a/network/network_utils.py +++ b/network/network_utils.py @@ -12,7 +12,7 @@ import numpy as np class BaseNet: def set_gpu(self, gpu): if gpu >= 0 and torch.cuda.is_available(): - self.device = torch.device('gpu:%d' % (gpu)) + self.device = torch.device('cuda:%d' % (gpu)) else: self.device = torch.device('cpu') self.to(self.device)