remove __ost_init__

NVIDIA · Oct 25, 2024 · fadca1e · fadca1e
1 parent bb1c874
commit fadca1e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py b/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py
@@ -23,8 +23,9 @@
 class MegatronDataModule(pl.LightningDataModule):
     """A mixin that adds a `state_dict` and `load_state_dict` method for datamodule training resumption in NeMo."""
 
-    def __post_init__(self):
+    def _init__(self, *args, **kwargs):
         """Set init_global_step to 0 for datamodule resumption."""
+        super().__init__(self, *args, **kwargs)
         self.init_global_step = 0
 
     def update_init_global_step(self):