(self, module, device_ids=None, output_device=None, dim=0, gather=True)
| 154 | # TODO: update notes/cuda.rst when this class handles 8+ GPUs well |
| 155 | |
| 156 | def __init__(self, module, device_ids=None, output_device=None, dim=0, gather=True): |
| 157 | super(MyDataParallel, self).__init__() |
| 158 | |
| 159 | if not torch.cuda.is_available(): |
| 160 | self.module = module |
| 161 | self.device_ids = [] |
| 162 | return |
| 163 | |
| 164 | if device_ids is None: |
| 165 | device_ids = list(range(torch.cuda.device_count())) |
| 166 | if output_device is None: |
| 167 | output_device = device_ids[0] |
| 168 | self.dim = dim |
| 169 | self.module = module |
| 170 | self.device_ids = device_ids |
| 171 | self.output_device = output_device |
| 172 | self.gather_bool = gather |
| 173 | |
| 174 | _check_balance(self.device_ids) |
| 175 | |
| 176 | if len(self.device_ids) == 1: |
| 177 | self.module.cuda(device_ids[0]) |
| 178 | |
| 179 | def forward(self, *inputs, **kwargs): |
| 180 | if not self.device_ids: |
nothing calls this directly
no test coverage detected