switch to zeros param init

NVIDIA · Jan 21, 2025 · 7b69b9e · 7b69b9e
1 parent 45ebc52
commit 7b69b9e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/model/model.py b/sub-packages/bionemo-esm2/src/bionemo/esm2/model/model.py
@@ -362,3 +362,6 @@ class ESM2Config(ESM2GenericConfig, iom.IOMixinWithGettersSetters):
     model_cls: Type[ESM2Model] = ESM2Model
     num_layers: int = 33  # 650M
     hidden_size: int = 1280  # 650M
+    output_layer_init_method: Callable = (
+        torch.nn.init.zeros_
+    )  # TODO make param init reproducible; remove after debugging