Fix monkeypatching of _FabricModule methods

src/lightning/fabric/CHANGELOG.md

@@ -49,7 +49,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed an issue causing a TypeError when using `torch.compile` as a decorator ([#19627](https://github.com/Lightning-AI/pytorch-lightning/pull/19627))
 
--
+- Fixed issue where some model methods couldn't be monkeypatched after being Fabric wrapped ([#19705](https://github.com/Lightning-AI/pytorch-lightning/pull/19705))
 
 -
 

src/lightning/fabric/wrappers.py

@@ -263,7 +263,7 @@ def __setattr__(self, name: str, value: Any) -> None:
         original_has_attr = hasattr(original_module, name)
         # Can't use super().__getattr__ because nn.Module only checks _parameters, _buffers, and _modules
         # Can't use self.__getattr__ because it would pass through to the original module
-        fabric_has_attr = name in self.__dict__
+        fabric_has_attr = hasattr(self, name)
 
         if not (original_has_attr or fabric_has_attr):
             setattr(original_module, name, value)

tests/tests_fabric/test_wrappers.py

@@ -170,6 +170,13 @@ def __init__(self):
     assert linear in fabric_module.modules()
     assert linear in original_module.modules()
 
+    # Check monkeypatching of methods
+    model = _FabricModule(Mock(), Mock())
+    assert isinstance(model, _FabricModule)
+    original = id(model.forward)
+    model.forward = lambda *_: None
+    assert id(model.forward) != original
+
 
 def test_fabric_module_state_dict_access():
     """Test that state_dict access passes through to the original module."""