(cls, path: str | Path, device: str = "cpu")
| 39 | |
| 40 | @classmethod |
| 41 | def load(cls, path: str | Path, device: str = "cpu") -> RouterCheckpoint: |
| 42 | state = torch.load(path, map_location=device, weights_only=True) |
| 43 | hidden_dim = state["hidden_dim"] |
| 44 | bottleneck_dim = state["bottleneck_dim"] |
| 45 | routers = {} |
| 46 | for layer_idx_str, router_state in state["router_layers"].items(): |
| 47 | layer_idx = int(layer_idx_str) if isinstance(layer_idx_str, str) else layer_idx_str |
| 48 | router = TokenRouter(hidden_dim, bottleneck_dim) |
| 49 | router.load_state_dict(router_state) |
| 50 | router.to(device) |
| 51 | routers[layer_idx] = router |
| 52 | return cls(routers=routers, hidden_dim=hidden_dim, bottleneck_dim=bottleneck_dim) |
nothing calls this directly
no test coverage detected