From ea47c120aa60f121d08ecb3c88703161ef029687 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 23 Oct 2026 12:33:49 -0600 Subject: [PATCH] Fix issue of saving full models when block compile is on --- toolkit/models/v2/_mixin.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 9f699fa..bf507b0 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -903,6 +903,8 @@ class OstrisModelMixin: value = value.to(dtype=dtype) state_dict[key] = value.detach().to("cpu").contiguous() state_dict.update(q_entries) + # block_compile wraps blocks in OptimizedModule: drop its "._orig_mod." key segment + state_dict = {k.replace("._orig_mod.", "."): v for k, v in state_dict.items()} state_dict = self.convert_state_dict_on_save(state_dict) parent = os.path.dirname(output_path)