(self, *inputs: Any, **kwargs: Any)
| 162 | self.module.to(self.src_device_obj) |
| 163 | |
| 164 | def forward(self, *inputs: Any, **kwargs: Any) -> Any: |
| 165 | with torch.autograd.profiler.record_function("DataParallel.forward"): |
| 166 | if not self.device_ids: |
| 167 | return self.module(*inputs, **kwargs) |
| 168 | |
| 169 | for t in chain(self.module.parameters(), self.module.buffers()): |
| 170 | if t.device != self.src_device_obj: |
| 171 | raise RuntimeError("module must have its parameters and buffers " |
| 172 | f"on device {self.src_device_obj} (device_ids[0]) but found one of " |
| 173 | f"them on device: {t.device}") |
| 174 | |
| 175 | inputs, module_kwargs = self.scatter(inputs, kwargs, self.device_ids) |
| 176 | # for forward function without any inputs, empty list and dict will be created |
| 177 | # so the module can be executed on one device which is the first one in device_ids |
| 178 | if not inputs and not module_kwargs: |
| 179 | inputs = ((),) |
| 180 | module_kwargs = ({},) |
| 181 | |
| 182 | if len(self.device_ids) == 1: |
| 183 | return self.module(*inputs[0], **module_kwargs[0]) |
| 184 | replicas = self.replicate(self.module, self.device_ids[:len(inputs)]) |
| 185 | outputs = self.parallel_apply(replicas, inputs, module_kwargs) |
| 186 | return self.gather(outputs, self.output_device) |
| 187 | |
| 188 | def replicate(self, module: T, device_ids: Sequence[Union[int, torch.device]]) -> List[T]: |
| 189 | return replicate(module, device_ids, not torch.is_grad_enabled()) |
nothing calls this directly
no test coverage detected