(device, numId=0)
| 141 | |
| 142 | |
| 143 | def set_device(device, numId=0): |
| 144 | import torch |
| 145 | if device == 'gpu' and torch.cuda.is_available(): |
| 146 | device = torch.device(f'cuda:{numId}') |
| 147 | else: |
| 148 | logger.info('GPU is not available, using CPU.') |
| 149 | device = torch.device('cpu') |
| 150 | return device |
| 151 | |
| 152 | |
| 153 | class OpenDetector(object): |