(self, method_name, *args, **kwargs)
| 33 | return func(*args, **kwargs) |
| 34 | |
| 35 | def execute_module(self, method_name, *args, **kwargs): |
| 36 | print(f"execute_module called with method={method_name}") |
| 37 | func = getattr(self._engine.module, method_name) |
| 38 | return func(*args, **kwargs) |
| 39 | |
| 40 | def get_model_size_on_rank_zero(self): |
| 41 | import torch |