fix bug on infer : device select and repvgg convert

furyhawk · Jul 29, 2022 · fb94337 · fb94337
1 parent b33352c
commit fb94337
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/yolov6/core/inferer.py b/yolov6/core/inferer.py
@@ -29,7 +29,7 @@ def __init__(self, source, weights, device, yaml, img_size, half):
         self.device = device
         self.img_size = img_size
         cuda = self.device != 'cpu' and torch.cuda.is_available()
-        self.device = torch.device('cuda:0' if cuda else 'cpu')
+        self.device = torch.device(f'cuda:{device}' if cuda else 'cpu')
         self.model = DetectBackend(weights, device=self.device)
         self.stride = self.model.stride
         self.class_names = load_yaml(yaml)['names']
@@ -49,7 +49,7 @@ def __init__(self, source, weights, device, yaml, img_size, half):
         self.files = LoadData(source)
 
         # Switch model to deploy status
-        self.model_switch(self.model, self.img_size)
+        self.model_switch(self.model.model, self.img_size)
 
     def model_switch(self, model, img_size):
         ''' Model switch to deploy status '''