Commit 8e66a12f authored by Haoran Li's avatar Haoran Li Committed by Facebook Github Bot
Browse files

fixes for exporter issue of bi-transformer model (#597) 

Summary:
Pull Request resolved: https://github.com/pytorch/fairseq/pull/597

Pull Request resolved: https://github.com/facebookresearch/pytext/pull/424

Fixes two issues:
1. the new Layernorm has issues in exporting
2. fix tensorboard writing by using the "RAW" operator_export_type

Differential Revision: D14610694

fbshipit-source-id: 1b859f54c571a90766128ab28539a9901375c3e6
parent b26d6b58

fairseq/modules/layer_norm.py

+2 −2
Original line number Diff line number Diff line
@@ -8,8 +8,8 @@ 
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)