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())