Model to device. It also warps models with DistributedDataParallel or DataParallel. Args: net (nn.Module)
(self, net)
| 86 | return self.log_dict |
| 87 | |
| 88 | def model_to_device(self, net): |
| 89 | """Model to device. It also warps models with DistributedDataParallel |
| 90 | or DataParallel. |
| 91 | |
| 92 | Args: |
| 93 | net (nn.Module) |
| 94 | """ |
| 95 | net = net.to(self.device) |
| 96 | if self.opt['dist']: |
| 97 | find_unused_parameters = self.opt.get('find_unused_parameters', False) |
| 98 | net = DistributedDataParallel( |
| 99 | net, device_ids=[torch.cuda.current_device()], find_unused_parameters=find_unused_parameters) |
| 100 | elif self.opt['num_gpu'] > 1: |
| 101 | net = DataParallel(net) |
| 102 | return net |
| 103 | |
| 104 | def get_optimizer(self, optim_type, params, lr, **kwargs): |
| 105 | if optim_type == 'Adam': |