From d705961f50a449fa81160ee4059069d3be8f48b9 Mon Sep 17 00:00:00 2001 From: Micha Livne Date: Thu, 17 Feb 2022 20:07:18 +0200 Subject: [PATCH] 1. Added another guard. Signed-off-by: Micha Livne --- nemo/collections/nlp/modules/common/megatron/clip_grads.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/collections/nlp/modules/common/megatron/clip_grads.py b/nemo/collections/nlp/modules/common/megatron/clip_grads.py index eb8023b47f51..da44618db02e 100644 --- a/nemo/collections/nlp/modules/common/megatron/clip_grads.py +++ b/nemo/collections/nlp/modules/common/megatron/clip_grads.py @@ -14,13 +14,13 @@ """Gradient clipping.""" -import amp_C import torch from torch._six import inf from nemo.collections.nlp.modules.common.megatron.module import param_is_not_shared try: + import amp_C from apex.multi_tensor_apply import multi_tensor_applier from apex.transformer import parallel_state from apex.transformer.tensor_parallel.layers import param_is_not_tensor_parallel_duplicate