Fixed an issue where tl_requires_grad wasn't being cleaned up properly.

johnmarktaylor91 · Aug 30, 2024 · 4da3e31 · 4da3e31
1 parent d01a3d1
commit 4da3e31
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torchlens/model_history.py b/torchlens/model_history.py
@@ -1591,8 +1591,9 @@ def restore_model_attributes(
             )
 
         for param in model.parameters():
-            param.requires_grad = getattr(param, "tl_requires_grad")
-            delattr(param, "tl_requires_grad")
+            if hasattr(param, "tl_requires_grad"):
+                param.requires_grad = getattr(param, "tl_requires_grad")
+                delattr(param, "tl_requires_grad")
 
     def undecorate_model_tensors(self, model: nn.Module):
         """Goes through a model and all its submodules, and unmutates any tensor attributes. Normally just clearing