| 5 | |
| 6 | |
| 7 | def set_lcd_model(cfg, logger=None, neptune=None): |
| 8 | #* set up GPU usage |
| 9 | use_cuda = torch.cuda.is_available() |
| 10 | gpu_ids = [] |
| 11 | if use_cuda: |
| 12 | gpu_c = torch.cuda.device_count() |
| 13 | if cfg.TRAINING.GPU.IDS is not None: |
| 14 | if len(cfg.TRAINING.GPU.IDS) <= gpu_c: |
| 15 | gpu_ids = cfg.TRAINING.GPU.IDS |
| 16 | else: |
| 17 | raise ValueError("Incorrect GPU Index, Please Check!") |
| 18 | else: |
| 19 | gpu_ids = np.arange(gpu_c).tolist() |
| 20 | cuda = "cuda" |
| 21 | device = torch.device(cuda if use_cuda else 'cpu') |
| 22 | lcd = LcdNet(cfg, logger, use_cuda, device, neptune, gpu_ids) |
| 23 | return lcd, [use_cuda, device, gpu_ids] |