Make FusedLayerNorm wrapper pickleable
Summary: See
import torch
-def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False):
- if not export and torch.cuda.is_available():
try:
- from apex.normalization import FusedLayerNorm
+ from apex.normalization import FusedLayerNorm as _FusedLayerNorm
+ has_fused_layernorm = True
- class _FusedLayerNorm(FusedLayerNorm):
+ class FusedLayerNorm(_FusedLayerNorm):
@torch.jit.unused
def forward(self, x):
return super().forward(x)
- return _FusedLayerNorm(normalized_shape, eps, elementwise_affine)
except ImportError:
- pass
+ has_fused_layernorm = False
+
+
+def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False):
+ if not export and torch.cuda.is_available() and has_fused_layernorm:
+ return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)
