(self, device)
| 41 | return output_names |
| 42 | |
| 43 | def set_device(self, device): |
| 44 | if isinstance(device, str): |
| 45 | device = torch.device(device) |
| 46 | if device.type == "cuda": |
| 47 | device = "{}:{}".format(device.type, device.index) |
| 48 | providers = [("CUDAExecutionProvider", {"device_id": device[-1:] if device[-1] in [str(_i) for _i in range(10)] else "0"}), "CPUExecutionProvider"] |
| 49 | else: |
| 50 | providers = ["CPUExecutionProvider"] |
| 51 | self.session.set_providers(providers) |
| 52 | self.device = device |
| 53 | |
| 54 | |
| 55 | class Yolo(SimpleOnnxInference): |