(self, model)
| 825 | return model |
| 826 | |
| 827 | def to_parallel(self, model) -> Union[nn.Module, TorchModel]: |
| 828 | # config format to reserve custom ddp |
| 829 | if self.cfg.get('parallel', None) is not None: |
| 830 | dp_cfg = deepcopy(self.cfg['parallel']) |
| 831 | dp_cfg.update( |
| 832 | dict(module=model, device_ids=[torch.cuda.current_device()])) |
| 833 | return build_parallel(dp_cfg) |
| 834 | |
| 835 | dp_cfg = dict( |
| 836 | type='DistributedDataParallel', |
| 837 | module=model, |
| 838 | find_unused_parameters=True, |
| 839 | device_ids=[torch.cuda.current_device()], |
| 840 | process_group=self.dp_group) |
| 841 | |
| 842 | return build_parallel(dp_cfg) |
| 843 | |
| 844 | def unwrap_module(self, model) -> Union[nn.Module, TorchModel]: |
| 845 | """Unwrap the model until it's a naked nn.Module. |
no test coverage detected