ci imports fix

Signed-off-by: dimapihtar <[email protected]>

nemo/collections/nlp/parts/nlp_overrides.py

@@ -60,14 +60,14 @@
 from nemo.collections.nlp.parts import utils_funcs
 from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
 from nemo.core.optim import MainParamsOptimizerWrapper
-from nemo.core.optim.distributed_adam import MegatronDistributedFusedAdam
 from nemo.core.optim.optimizers import init_optimizer_states
 from nemo.utils import AppState, logging
 from nemo.utils.get_rank import is_global_rank_zero
 from nemo.utils.model_utils import ckpt_to_dir, inject_model_parallel_rank, uninject_model_parallel_rank
 
 try:
     from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+    from nemo.core.optim.distributed_adam import MegatronDistributedFusedAdam
 
     HAVE_APEX = True
 

nemo/core/optim/distributed_adam.py

@@ -24,18 +24,11 @@
 
 from nemo.utils import str_to_dtype
 
-try:
-    from apex.contrib.optimizers.distributed_fused_adam import (
-        DistributedFusedAdam,
-        _disable_pre_forward_hook,
-        _multi_tensor_copy,
-    )
-
-    HAVE_APEX = True
-
-except (ImportError, ModuleNotFoundError):
-
-    HAVE_APEX = False
+from apex.contrib.optimizers.distributed_fused_adam import (
+    DistributedFusedAdam,
+    _disable_pre_forward_hook,
+    _multi_tensor_copy,
+)
 
 # Check if Transformer Engine has FP8 tensor class
 HAVE_TE_FP8TENSOR = False