Minor update

2026-09-09 11:13:47 +08:00 · 2018-05-18 14:24:45 -06:00
parent dcb363067e
commit 93a5c3fc04
2 changed files with 5 additions and 4 deletions
@@ -12,6 +12,7 @@ import numpy as np
 class BaseNet:
    def set_gpu(self, gpu):
        if gpu >= 0 and torch.cuda.is_available():
+            gpu = gpu % torch.cuda.device_count()
            self.device = torch.device('cuda:%d' % (gpu))
        else:
            self.device = torch.device('cpu')