|
|
@@ -26,8 +26,6 @@ class DefaultFinetuner(mixins._ModelMixin, mixins._DatasetMixin, mixins._Trainer
|
|
|
super(DefaultFinetuner, self).__init__(opts=opts, *args, **kwargs)
|
|
|
|
|
|
self.gpu_config(opts)
|
|
|
- cuda.get_device_from_id(self.device).use()
|
|
|
-
|
|
|
self.read_annotations(opts)
|
|
|
|
|
|
self.init_model(opts)
|
|
|
@@ -43,7 +41,11 @@ class DefaultFinetuner(mixins._ModelMixin, mixins._DatasetMixin, mixins._Trainer
|
|
|
|
|
|
def gpu_config(self, opts):
|
|
|
if -1 in opts.gpu:
|
|
|
- self.device = -1
|
|
|
+ self.device_id = -1
|
|
|
else:
|
|
|
- self.device = opts.gpu[0]
|
|
|
+ self.device_id = opts.gpu[0]
|
|
|
+
|
|
|
+ self.device = cuda.get_device_from_id(self.device_id)
|
|
|
+ self.device.use()
|
|
|
+ return self.device
|
|
|
|
