Fused Layer Norm not working out-of-the-box with upstream AMP

Error when using fusedlayernorm with upstream amp. pasting here the log I get from Bobby Chen:

  File "/workspace/electra/modeling_utils.py", line 1307, in forward
    hidden_states = self.LayerNorm(hidden_states + input_tensor)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 577, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/apex/normalization/fused_layer_norm.py", line 159, in forward
    input, self.weight, self.bias, self.normalized_shape,self.eps)
  File "/opt/conda/lib/python3.6/site-packages/apex/normalization/fused_layer_norm.py", line 25, in forward
    input_, ctx.normalized_shape, weight_, bias_, ctx.eps)
RuntimeError: expected scalar type Half but found Float

I think this is due to issue described here, and fused layer norm need to be updated
https://pytorch.org/docs/stable/notes/amp_examples.html#amp-custom-examples

@mcarilli @ptrblck to confirm

Start with apex/normalization/fused_layer_norm.py, especially the forward path shown in the traceback, and review the upstream AMP custom examples linked in the issue. Reproduce the fused layer norm call from modeling_utils.py with upstream AMP; done means the operation no longer fails with the Half-versus-Float scalar type error.