(
fabric: L.Fabric, model: nn.Module, checkpoint_path: Path, strict: bool = True
)
| 358 | |
| 359 | |
| 360 | def load_checkpoint( |
| 361 | fabric: L.Fabric, model: nn.Module, checkpoint_path: Path, strict: bool = True |
| 362 | ) -> None: |
| 363 | if isinstance(fabric.strategy, FSDPStrategy): |
| 364 | fabric.load_raw(checkpoint_path, model, strict=strict) |
| 365 | else: |
| 366 | state_dict = lazy_load(checkpoint_path) |
| 367 | state_dict = state_dict.get('model', state_dict) |
| 368 | model.load_state_dict(state_dict, strict=strict) |
| 369 | |
| 370 | |
| 371 | def flops_per_param( |