(self, *args, **kwargs)
| 174 | self._memory_controller = memory_controller |
| 175 | |
| 176 | def forward(self, *args, **kwargs): |
| 177 | if self._memory_controller is None or not torch.is_grad_enabled() or not self.training: |
| 178 | _, ret = self._forward_with_mem_ratio(*args, **kwargs) |
| 179 | else: |
| 180 | input_size = self._get_input_size(*args, **kwargs) |
| 181 | mem_ratio = self._memory_controller.get_mem_ratio(input_size) |
| 182 | mem_ratio, ret = self._forward_with_mem_ratio(*args, mem_ratio=mem_ratio, **kwargs) |
| 183 | self._memory_controller.update_run_states(input_size, mem_ratio) |
| 184 | return ret |
| 185 | |
| 186 | |
| 187 | class ElasticModuleMixin: |
no test coverage detected