Include initialize as config parameter

funkelab · Oct 12, 2023 · 406a8a2 · 406a8a2
1 parent 4b17700
commit 406a8a2
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/cellulus/configs/model_config.py b/cellulus/configs/model_config.py
@@ -40,6 +40,9 @@ class ModelConfig:
             checkpoint will be used to resume training, otherwise the network
             will be trained from scratch.
 
+        initialize (default: False)
+
+            If True, initialize the model weights with Kaiming Normal
 
     """
 
@@ -52,3 +55,4 @@ class ModelConfig:
         ]
     )
     checkpoint: Path = attrs.field(default=None, converter=to_path)
+    initialize: bool = attrs.field(default=False, validator=instance_of(bool))