diff --git a/src/transformers/models/bart/modeling_bart.py b/src/transformers/models/bart/modeling_bart.py index f631736c..7f4af885 100644 --- a/src/transformers/models/bart/modeling_bart.py +++ b/src/transformers/models/bart/modeling_bart.py @@ -110,13 +110,12 @@ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] def BartLayerNorm(normalized_shape: torch.Size, eps: float = 1e-5, elementwise_affine: bool = True): - if torch.cuda.is_available(): - try: - from apex.normalization import FusedLayerNorm + try: + from apex.normalization import FusedLayerNorm - return FusedLayerNorm(normalized_shape, eps, elementwise_affine) - except ImportError: - pass + return FusedLayerNorm(normalized_shape, eps, elementwise_affine) + except ImportError: + pass return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine) diff --git a/src/transformers/models/fsmt/modeling_fsmt.py b/src/transformers/models/fsmt/modeling_fsmt.py index 682d6af0..0cd07ed6 100644 --- a/src/transformers/models/fsmt/modeling_fsmt.py +++ b/src/transformers/models/fsmt/modeling_fsmt.py @@ -265,14 +265,12 @@ FSMT_INPUTS_DOCSTRING = r""" have_fused_layer_norm = False -if torch.cuda.is_available(): - try: - from apex.normalization import FusedLayerNorm - - have_fused_layer_norm = True - except ImportError: - pass +try: + from apex.normalization import FusedLayerNorm + have_fused_layer_norm = True +except ImportError: + pass LayerNorm = FusedLayerNorm if have_fused_layer_norm else torch.nn.LayerNorm diff --git a/src/transformers/models/prophetnet/modeling_prophetnet.py b/src/transformers/models/prophetnet/modeling_prophetnet.py index ffa669d6..17db02b5 100644 --- a/src/transformers/models/prophetnet/modeling_prophetnet.py +++ b/src/transformers/models/prophetnet/modeling_prophetnet.py @@ -511,13 +511,12 @@ class ProphetNetDecoderLMOutput(ModelOutput): def ProphetNetLayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True): - if torch.cuda.is_available(): - try: - from apex.normalization import FusedLayerNorm + try: + from apex.normalization import FusedLayerNorm - return FusedLayerNorm(normalized_shape, eps, elementwise_affine) - except ImportError: - pass + return FusedLayerNorm(normalized_shape, eps, elementwise_affine) + except ImportError: + pass return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)