(self, *args, **kwargs)
| 571 | lazy_stop_callback=None): |
| 572 | |
| 573 | def _step(self, *args, **kwargs): |
| 574 | for name, param in module.named_parameters(): |
| 575 | io_json[name] = { |
| 576 | 'data': { |
| 577 | 'sum': |
| 578 | numpify_tensor_nested( |
| 579 | detach_tensor_nested(param.data), reduction='sum'), |
| 580 | 'mean': |
| 581 | numpify_tensor_nested( |
| 582 | detach_tensor_nested(param.data), reduction='mean'), |
| 583 | }, |
| 584 | 'grad': { |
| 585 | 'sum': |
| 586 | numpify_tensor_nested( |
| 587 | detach_tensor_nested(param.grad), reduction='sum'), |
| 588 | 'mean': |
| 589 | numpify_tensor_nested( |
| 590 | detach_tensor_nested(param.grad), reduction='mean'), |
| 591 | } |
| 592 | } |
| 593 | ret = self.step_origin(*args, **kwargs) |
| 594 | for name, param in module.named_parameters(): |
| 595 | io_json[name]['data_after'] = { |
| 596 | 'sum': |
| 597 | numpify_tensor_nested( |
| 598 | detach_tensor_nested(param.data), reduction='sum'), |
| 599 | 'mean': |
| 600 | numpify_tensor_nested( |
| 601 | detach_tensor_nested(param.data), reduction='mean'), |
| 602 | } |
| 603 | if lazy_stop_callback is not None: |
| 604 | lazy_stop_callback() |
| 605 | return ret |
| 606 | |
| 607 | if not restore and not hasattr(optimizer, 'step_origin'): |
| 608 | optimizer.step_origin, optimizer.step = optimizer.step, type( |
nothing calls this directly
no test coverage detected
searching dependent graphs…