fix cherry-pick

Signed-off-by: ashors1 <ashors@nvidia.com>
NVIDIA · Jul 15, 2024 · 81302ba · 81302ba
1 parent 8141139
commit 81302ba
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/nemo/lightning/megatron_parallel.py b/nemo/lightning/megatron_parallel.py
@@ -129,7 +129,6 @@ def __init__(
         cpu: bool = False,
         convert_module_fn: Optional[Callable[[ModelT], nn.Module]] = None,
     ) -> None:
-        from apex.transformer.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes
         from megatron.core import parallel_state
         from megatron.core.tensor_parallel import set_defaults_if_not_set_tensor_model_parallel_attributes