(flags)
| 96 | |
| 97 | |
| 98 | def build_device(flags): |
| 99 | if flags.Global.use_gpu: |
| 100 | os.environ["CUDA_VISIBLE_DEVICES"] = flags.Global.gpu_num |
| 101 | if torch.cuda.is_available(): |
| 102 | device = torch.device(flags.Global.device) |
| 103 | gpu_count = torch.cuda.device_count() |
| 104 | else: |
| 105 | device = torch.device("cpu") |
| 106 | gpu_count = 0 |
| 107 | else: |
| 108 | device = torch.device("cpu") |
| 109 | gpu_count = 0 |
| 110 | return device, gpu_count |
| 111 | |
| 112 | |
| 113 | def build_loss(flags): |