late import prettytable (NVIDIA#10912)

Signed-off-by: Maanu Grover <[email protected]>
artbataev · Oct 22, 2024 · 23eee74 · 23eee74
1 parent d9d1f90
commit 23eee74
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/lightning/pytorch/callbacks/debugging.py b/nemo/lightning/pytorch/callbacks/debugging.py
@@ -2,7 +2,6 @@
 
 import pytorch_lightning as pl
 import torch
-from prettytable import PrettyTable
 from pytorch_lightning.callbacks import Callback
 
 from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule
@@ -141,6 +140,8 @@ def find_grad_tensor(param: torch.Tensor) -> Optional[torch.Tensor]:
 
         # create table only if there is something to print
         if any(param_keys) or any(grad_keys):
+            from prettytable import PrettyTable
+
             debug_table = PrettyTable()
             debug_table.add_column("Parameter", names_col)