From 71dbb4e6de4aadad6c3e25923e46abf9acb50ea4 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Thu, 11 Sep 2025 17:35:25 -0700 Subject: [PATCH] fix state dict loading Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags: --- torchtitan/components/checkpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchtitan/components/checkpoint.py b/torchtitan/components/checkpoint.py index e9e7014425..14a28eb206 100644 --- a/torchtitan/components/checkpoint.py +++ b/torchtitan/components/checkpoint.py @@ -70,7 +70,7 @@ def _get_state_dict(self) -> dict[str, Any]: return state_dict def state_dict(self) -> dict[str, Any]: - return self.cache_state_dict + return self._get_state_dict() def load_state_dict(self, state_dict: dict[str, Any]) -> None: func = functools.partial(