fix get_linear_nonlinear_params

Lightning-AI · May 6, 2024 · bd73193 · bd73193
1 parent fc17c57
commit bd73193
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/litgpt/utils.py b/litgpt/utils.py
@@ -490,13 +490,12 @@ def choose_logger(
 def get_linear_nonlinear_params(model):
     linear_params = []
     nonlinear_params = []
-
     for module in model.modules():
         if isinstance(module, torch.nn.Linear):
-            linear_params.extend(list(model.parameters()))
+            linear_params.extend(list(module.parameters()))
         else:
-            nonlinear_params.extend(list(model.parameters()))
 
-    # Make extra sure that there is no overlap
-    linear_params = list(set(linear_params) - set(nonlinear_params))
+            nonlinear_params.extend(list(module.parameters()))
+    linear_params = list(set(linear_params))
+    nonlinear_params = list(set(nonlinear_params) - set(linear_params))
     return linear_params, nonlinear_params