[PyTorch] Fix bug when reshaping norm output from LayerNormLinear (#1741)

Fix bug when reshaping norm output from LayerNormLinear
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/module/layernorm_linear.py

@@ -467,7 +467,7 @@ class _LayerNormLinear(torch.autograd.Function):
                 shape = list(inp_shape)
                 shape[0] *= tp_size
                 return out, ln_out_return.view(shape)
-            return out, ln_out_return.view_as(inp)
+            return out, ln_out_return.view(inp_shape)
         return out
 
     @staticmethod

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -612,7 +612,7 @@ class _LayerNormMLP(torch.autograd.Function):
                 shape = list(inp_shape)
                 shape[0] *= tp_size
                 return fc2_out, ln_out_return.view(shape)
-            return fc2_out, ln_out_return.view_as(inp)
+            return fc2_out, ln_out_return.view(inp_shape)
         return fc2_out
 
     @staticmethod