fix recursion error when setting tp_wrapped_module #122

BlackSamorez · Oct 2, 2023 · 14bdd34 · 14bdd34
1 parent a7d1939
commit 14bdd34
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/tensor_parallel/wrapper.py b/src/tensor_parallel/wrapper.py
@@ -73,3 +73,8 @@ def forward(self, *args, **kwargs):
 
     def __getattr__(self, attr):
         return getattr(self.tp_wrapped_module, attr)
+
+    def __setattr__(self, attr, value):
+        super().__setattr__(attr, value)
+        if attr == "tp_wrapped_module":
+            self.__dict__["tp_wrapped_module"] = value  # to access without getattr, nn.Module removed it from __dict__