(self, *args, **kwargs)
| 217 | self._memory_controller = memory_controller |
| 218 | |
| 219 | def forward(self, *args, **kwargs): |
| 220 | if self._memory_controller is None or not torch.is_grad_enabled() or not self.training: |
| 221 | ret = super().forward(*args, **kwargs) |
| 222 | else: |
| 223 | input_size = self._get_input_size(*args, **kwargs) |
| 224 | mem_ratio = self._memory_controller.get_mem_ratio(input_size) |
| 225 | with self.with_mem_ratio(mem_ratio) as exact_mem_ratio: |
| 226 | ret = super().forward(*args, **kwargs) |
| 227 | self._memory_controller.update_run_states(input_size, exact_mem_ratio) |
| 228 | return ret |
nothing calls this directly
no test coverage detected