diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/model/lr_scheduler.py b/sub-packages/bionemo-llm/src/bionemo/llm/model/lr_scheduler.py
similarity index 100%
rename from sub-packages/bionemo-esm2/src/bionemo/esm2/model/lr_scheduler.py
rename to sub-packages/bionemo-llm/src/bionemo/llm/model/lr_scheduler.py
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/train.py b/sub-packages/bionemo-llm/src/bionemo/llm/train.py
index 084513aad0..a3f7ceae85 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/train.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/train.py
@@ -126,13 +126,6 @@ def setup_trainer(
             )
         )
 
-    # TODO set these as flags, the following are needed:
-    """
-    nsys_profiling (bool)
-    nsys_start_step (int) when to start profiling
-    nsys_end_step (int) when to stop profiling
-    nsys_ranks (List[int]) which ranks to profile.
-    """
     if nsys_config:
         if nsys_config.end_step is None:
             nsys_config.end_step = training_config.max_steps