FIX #124 trial 1: set_layer_trainable

MannLabs · Dec 31, 2023 · 8b85dc6 · 8b85dc6
1 parent a4a98d2
commit 8b85dc6
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/peptdeep/model/model_interface.py b/peptdeep/model/model_interface.py
@@ -256,15 +256,13 @@ def set_bert_trainable(self,
                 for layer in bert_layer_idxes
             ]
         )
-    
+
     def set_layer_trainable(self, 
         trainable=False, 
         layer_names=[],
     ):
         for layer in layer_names:
-            self.model.get_submodule(
-                layer
-            ).requires_grad_(trainable)
+            self.model.get_submodule(layer).requires_grad_(trainable)
 
     def train_with_warmup(self,
         precursor_df: pd.DataFrame,