Figure out which device (i.e., gpu or cpu) to use.
(process_info: T.Dict[str, T.Any])
| 598 | |
| 599 | @staticmethod |
| 600 | def _determine_device(process_info: T.Dict[str, T.Any]): |
| 601 | """Figure out which device (i.e., gpu or cpu) to use.""" |
| 602 | if torch.cuda.is_available() and process_info["n_gpus"] > 0: |
| 603 | gpu_id = process_info["rank"] % process_info["n_gpus"] |
| 604 | assert ( |
| 605 | torch.cuda.device_count() >= process_info["n_gpus"] |
| 606 | ), f'{torch.cuda.device_count()} {process_info["n_gpus"]}' |
| 607 | # set the default cuda device |
| 608 | torch.cuda.set_device(gpu_id) |
| 609 | device = torch.device(f"cuda:{gpu_id}") |
| 610 | else: |
| 611 | device = torch.device("cpu") |
| 612 | return device |
| 613 | |
| 614 | def _send_models_to_device(self, device: torch.device): |
| 615 | """Send all base model and nn.modules to device.""" |