From 3b4b1ad2e548bfc210c4e870ec14144f12a155bb Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 26 Jul 2026 13:07:53 -0600 Subject: [PATCH] Ensure force_patch_weights is set when partially loading/unloading models --- nodes.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/nodes.py b/nodes.py index 4683514..15e9176 100644 --- a/nodes.py +++ b/nodes.py @@ -89,6 +89,12 @@ def pin_weight_to_device(self, key): mmap_released = False named_modules_to_munmap = {} + def partially_unload(self, *args, force_patch_weights=False, **kwargs): + return super().partially_unload(*args, force_patch_weights=True, **kwargs) + + def partially_load(self, *args, force_patch_weights=False, **kwargs): + return super().partially_load(*args, force_patch_weights=True, **kwargs) + def load(self, *args, force_patch_weights=False, **kwargs): if not self.mmap_released: self.named_modules_to_munmap = dict(self.model.named_modules())