Fix interaction between tracing and FusedLayerNorm
Summary: Pull Request resolved:
@@ -10,7 +10,14 @@ 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
- return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
+
+ class _FusedLayerNorm(FusedLayerNorm):
+
+ @torch.jit.unused
+ def forward(self, x):
+ return super().forward(x)
+
+ return _FusedLayerNorm(normalized_shape, eps, elementwise_affine)
except ImportError:
pass
return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)
