rename MegatronDataModule

NVIDIA · Oct 25, 2024 · bb1c874 · bb1c874
1 parent a05c230
commit bb1c874
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py b/sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py
@@ -32,7 +32,7 @@
 Mode = Literal["train", "validation", "test"]
 
 
-class ESMDataModule(MegatronDatamodule):
+class ESMDataModule(MegatronDataModule):
     """LightningDataModule wrapper of `ESMDataset`."""
 
     def __init__(

diff --git a/sub-packages/bionemo-geneformer/src/bionemo/geneformer/data/singlecell/datamodule.py b/sub-packages/bionemo-geneformer/src/bionemo/geneformer/data/singlecell/datamodule.py
@@ -39,7 +39,7 @@
 __all__: Sequence[str] = ("SingleCellDataModule",)
 
 
-class SingleCellDataModule(MegatronDatamodule):
+class SingleCellDataModule(MegatronDataModule):
     """LightningDataModule wrapper of `SingleCellDataset`
 
     Args:

diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py b/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py
@@ -20,7 +20,7 @@
 from nemo.utils import logging
 
 
-class MegatronDatamodule(pl.LightningDataModule):
+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):