Method save

python/TIDE/router.py:28–38

(self, path: str | Path)

    bottleneck_dim: int

    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        state = {
            "hidden_dim": self.hidden_dim,
            "bottleneck_dim": self.bottleneck_dim,
            "router_layers": {},
        }
        for layer_idx, router in self.routers.items():
            state["router_layers"][layer_idx] = router.state_dict()
        torch.save(state, path)

    @classmethod
    def load(cls, path: str | Path, device: str = "cpu") -> RouterCheckpoint: