Set export=True in LayerNorm when under TorchScript execution
Summary: When under TorchScript execution, set export=True as FusedLayerNorm doesn't work with JIT yet (See torch.jit.unused decorator).
@@ -27,6 +27,8 @@ except ImportError:
def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False):
+ if torch.jit.is_scripting():
+ export = True
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)
