fixes for exporter issue of bi-transformer model
Summary:
Pull Request resolved:
Pull Request resolved:
Fixes two issues:
1. the new Layernorm has issues in exporting
2. fix tensorboard writing by using the "RAW" operator_export_type
import torch
-def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True):
- if torch.cuda.is_available():
+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)
