Repository: pytorch
Issue: Detailed error: Tensor-likes are not close! When use torch.jit.trace
Labels: oncall: jit

### 🐛 Describe the bug

When run follows code:

```
import torch

decoder_layer = torch.nn.TransformerDecoderLayer(
    d_model=512, nhead=8, batch_first=True,
    dropout=0.0, dtype=torch.bfloat16
)
model = torch.nn.TransformerDecoder(decoder_layer, num_layers=1)

src = torch.rand((1, 10, 512), dtype=torch.bfloat16)
tgt = torch.rand((1, 20, 512), dtype=torch.bfloat16)

model_trace = flatten_trace(model, src, tgt) # my flatten trace function

model_jit = torch.jit.trace(model_trace, (src, tgt))

out_jit = model_jit(src, tgt)
```
An error occurred: 

TracerWarning: Output nr 1. of the traced function does not match the corresponding output of the Python function. Detailed error: Tensor-likes are not close!
Mismatched elements: 1607 / 5120 (31.4%)
Greatest absolute difference: 0.015625 at index (0, 0, 63) (up to 1e-05 allowed)
Greatest relative difference: 2.4893617021276597 at index (0, 0, 368) (up to 1e-05 allowed)

flatten_trace is my custom trace function, call this generate model:
```
def forward(self, tensor_0, tensor_1):
    with torch.autograd.profiler.record_function("FlattenModule"):
        layers_0_self_attn_in_proj_weight = self.layers_0_self_attn_in_proj_weight
        layers_0_self_attn_in_proj_bias = self.layers_0_self_attn_in_proj_bias
        layers_0_self_attn_out_proj_weight = self.layers_0_self_attn_out_proj_weight
        layers_0_self_attn_out_proj_bias = self.layers_0_self_attn_out_proj_bias
        layers_0_multihead_attn_in_proj_weight = self.layers_0_multihead_attn_in_proj_weight
        layers_0_multihead_attn_in_proj_bias = self.layers_0_multihead_attn_in_proj_bias
        layers_0_multihead_attn_out_proj_weight = self.layers_0_multihead_attn_out_proj_weight
        layers_0_multihead_attn_out_proj_bias = self.layers_0_multihead_attn_out_proj_bias
        layers_0_linear1_weight = self.layers_0_linear1_weight
        layers_0_linear1_bias = self.layers_0_linear1_bias
        layers_0_linear2_weight = self.layers_0_linear2_weight
        layers_0_linear2_bias = self.layers_0_linear2_bias
        layers_0_norm1_weight = self.layers_0_norm1_weight
        layers_0_norm1_bias = self.layers_0_norm1_bias
        layers_0_norm2_weight = self.layers_0_norm2_weight
        layers_0_norm2_bias = self.layers_0_norm2_bias
        layers_0_norm3_weight = self.layers_0_norm3_weight
        layers_0_norm3_bias = self.layers_0_norm3_bias
        transpose = torch.ops.aten.transpose(tensor_0, 1, 0)
        t = torch.ops.aten.t(layers_0_self_attn_in_proj_weight);  layers_0_self_attn_in_proj_weight = None
        view = torch.ops.aten.view(transpose, [10, 512]);  transpose = None
        addmm = torch.ops.aten.addmm(layers_0_self_attn_in_proj_bias, view, t);  layers_0_self_attn_in_proj_bias = view = t = None
        view_1 = torch.ops.aten.view(addmm, [10, 1, 1536]);  addmm = None
        split = torch.ops.aten.split(view_1, 512, -1);  view_1 = None
        getitem = split[0]
        getitem_1 = split[1]
        getitem_2 = split[2];  split = None
        clone = torch.ops.aten.clone(getitem, memory_format = torch.contiguous_format);  getitem = None
        view_2 = torch.ops.aten.view(clone, [10, 8, 64]);  clone = None
        transpose_1 = torch.ops.aten.transpose(view_2, 0, 1);  view_2 = None
        clone_1 = torch.ops.aten.clone(getitem_1, memory_format = torch.contiguous_format);  getitem_1 = None
        view_3 = torch.ops.aten.view(clone_1, [10, 8, 64]);  clone_1 = None
        transpose_2 = torch.ops.aten.transpose(view_3, 0, 1);  view_3 = None
        clone_2 = torch.ops.aten.clone(getitem_2, memory_format = torch.contiguous_format);  getitem_2 = None
        view_4 = torch.ops.aten.view(clone_2, [10, 8, 64]);  clone_2 = None
        transpose_3 = torch.ops.aten.transpose(view_4, 0, 1);  view_4 = None
        div = torch.ops.aten.div(transpose_1, 8.0);  transpose_1 = None
        transpose_4 = torch.ops.aten.transpose(transpose_2, -2, -1);  transpose_2 = None
        bmm = torch.ops.aten.bmm(div, transpose_4);  div = transpose_4 = None
        _softmax = torch.ops.aten._softmax(bmm, -1, False);  bmm = None
        bmm_1 = torch.ops.aten.bmm(_softmax, transpose_3);  _softmax = transpose_3 = None
        transpose_5 = torch.ops.aten.transpose(bmm_1, 0, 1);  bmm_1 = None
        clone_3 = torch.ops.aten.clone(transpose_5, memory_format = torch.contiguous_format);  transpose_5 = None
        view_5 = torch.ops.aten.view(clone_3, [10, 512]);  clone_3 = None
        t_1 = torch.ops.aten.t(layers_0_self_attn_out_proj_weight);  layers_0_self_attn_out_proj_weight = None
        addmm_1 = torch.ops.aten.addmm(layers_0_self_attn_out_proj_bias, view_5, t_1);  layers_0_self_attn_out_proj_bias = view_5 = t_1 = None
        view_6 = torch.ops.aten.view(addmm_1, [10, 1, 512]);  addmm_1 = None
        transpose_6 = torch.ops.aten.transpose(view_6, 1, 0);  view_6 = None
        add = torch.ops.aten.add(tensor_0, transpose_6);  tensor_0 = transpose_6 = None
        native_layer_norm = torch.ops.aten.native_layer_norm(add, [512], layers_0_norm1_weight, layers_0_norm1_bias, 1e-05);  add = layers_0_norm1_weight = layers_0_norm1_bias = None
        getitem_3 = native_layer_norm[0];  native_layer_norm = None
        transpose_7 = torch.ops.aten.transpose(getitem_3, 1, 0)
        transpose_8 = torch.ops.aten.transpose(tensor_1, 1, 0);  tensor_1 = None
        split_with_sizes = torch.ops.aten.split_with_sizes(layers_0_multihead_attn_in_proj_weight, [512, 1024]);  layers_0_multihead_attn_in_proj_weight = None
        getitem_6 = split_with_sizes[0]
        getitem_7 = split_with_sizes[1];  split_with_sizes = None
        split_with_sizes_1 = torch.ops.aten.split_with_sizes(layers_0_multihead_attn_in_proj_bias, [512, 1024]);  layers_0_multihead_attn_in_proj_bias = None
        getitem_8 = split_with_sizes_1[0]
        getitem_9 = split_with_sizes_1[1];  split_with_sizes_1 = None
        t_2 = torch.ops.aten.t(getitem_6);  getitem_6 = None
        view_7 = torch.ops.aten.view(transpose_7, [10, 512]);  transpose_7 = None
        addmm_2 = torch.ops.aten.addmm(getitem_8, view_7, t_2);  getitem_8 = view_7 = t_2 = None
        view_8 = torch.ops.aten.view(addmm_2, [10, 1, 512]);  addmm_2 = None
        t_3 = torch.ops.aten.t(getitem_7);  getitem_7 = None
        view_9 = torch.ops.aten.view(transpose_8, [20, 512]);  transpose_8 = None
        addmm_3 = torch.ops.aten.addmm(getitem_9, view_9, t_3);  getitem_9 = view_9 = t_3 = None
        view_10 = torch.ops.aten.view(addmm_3, [20, 1, 1024]);  addmm_3 = None
        split_1 = torch.ops.aten.split(view_10, 512, -1);  view_10 = None
        getitem_10 = split_1[0]
        getitem_11 = split_1[1];  split_1 = None
        view_11 = torch.ops.aten.view(view_8, [10, 8, 64]);  view_8 = None
        transpose_9 = torch.ops.aten.transpose(view_11, 0, 1);  view_11 = None
        clone_4 = torch.ops.aten.clone(getitem_10, memory_format = torch.contiguous_format);  getitem_10 = None
        view_12 = torch.ops.aten.view(clone_4, [20, 8, 64]);  clone_4 = None
        transpose_10 = torch.ops.aten.transpose(view_12, 0, 1);  view_12 = None
        clone_5 = torch.ops.aten.clone(getitem_11, memory_format = torch.contiguous_format);  getitem_11 = None
        view_13 = torch.ops.aten.view(clone_5, [20, 8, 64]);  clone_5 = None
        transpose_11 = torch.ops.aten.transpose(view_13, 0, 1);  view_13 = None
        div_1 = torch.ops.aten.div(transpose_9, 8.0);  transpose_9 = None
        transpose_12 = torch.ops.aten.transpose(transpose_10, -2, -1);  transpose_10 = None
        bmm_2 = torch.ops.aten.bmm(div_1, transpose_12);  div_1 = transpose_12 = None
        _softmax_1 = torch.ops.aten._softmax(bmm_2, -1, False);  bmm_2 = None
        bmm_3 = torch.ops.aten.bmm(_softmax_1, transpose_11);  _softmax_1 = transpose_11 = None
        transpose_13 = torch.ops.aten.transpose(bmm_3, 0, 1);  bmm_3 = None
        clone_6 = torch.ops.aten.clone(transpose_13, memory_format = torch.contiguous_format);  transpose_13 = None
        view_14 = torch.ops.aten.view(clone_6, [10, 512]);  clone_6 = None
        t_4 = torch.ops.aten.t(layers_0_multihead_attn_out_proj_weight);  layers_0_multihead_attn_out_proj_weight = None
        addmm_4 = torch.ops.aten.addmm(layers_0_multihead_attn_out_proj_bias, view_14, t_4);  layers_0_multihead_attn_out_proj_bias = view_14 = t_4 = None
        view_15 = torch.ops.aten.view(addmm_4, [10, 1, 512]);  addmm_4 = None
        transpose_14 = torch.ops.aten.transpose(view_15, 1, 0);  view_15 = None
        add_1 = torch.ops.aten.add(getitem_3, transpose_14);  getitem_3 = transpose_14 = None
        native_layer_norm_1 = torch.ops.aten.native_layer_norm(add_1, [512], layers_0_norm2_weight, layers_0_norm2_bias, 1e-05);  add_1 = layers_0_norm2_weight = layers_0_norm2_bias = None
        getitem_12 = native_layer_norm_1[0];  native_layer_norm_1 = None
        t_5 = torch.ops.aten.t(layers_0_linear1_weight);  layers_0_linear1_weight = None
        view_16 = torch.ops.aten.view(getitem_12, [10, 512])
        addmm_5 = torch.ops.aten.addmm(layers_0_linear1_bias, view_16, t_5);  layers_0_linear1_bias = view_16 = t_5 = None
        view_17 = torch.ops.aten.view(addmm_5, [1, 10, 2048]);  addmm_5 = None
        relu = torch.ops.aten.relu(view_17);  view_17 = None
        t_6 = torch.ops.aten.t(layers_0_linear2_weight);  layers_0_linear2_weight = None
        view_18 = torch.ops.aten.view(relu, [10, 2048]);  relu = None
        addmm_6 = torch.ops.aten.addmm(layers_0_linear2_bias, view_18, t_6);  layers_0_linear2_bias = view_18 = t_6 = None
        view_19 = torch.ops.aten.view(addmm_6, [1, 10, 512]);  addmm_6 = None
        add_2 = torch.ops.aten.add(getitem_12, view_19);  getitem_12 = view_19 = None
        native_layer_norm_2 = torch.ops.aten.native_layer_norm(add_2, [512], layers_0_norm3_weight, layers_0_norm3_bias, 1e-05);  add_2 = layers_0_norm3_weight = layers_0_norm3_bias = None
        getitem_15 = native_layer_norm_2[0];  native_layer_norm_2 = None
        return getitem_15
```

after call torch.jit.trace() function,  generate model:
```
def forward(self,
    tensor_0: Tensor,
    tensor_1: Tensor) -> Tensor:
  layers_0_norm3_bias = self.layers_0_norm3_bias
  layers_0_norm3_weight = self.layers_0_norm3_weight
  layers_0_linear2_bias = self.layers_0_linear2_bias
  layers_0_linear2_weight = self.layers_0_linear2_weight
  layers_0_linear1_bias = self.layers_0_linear1_bias
  layers_0_linear1_weight = self.layers_0_linear1_weight
  layers_0_norm2_bias = self.layers_0_norm2_bias
  layers_0_norm2_weight = self.layers_0_norm2_weight
  layers_0_multihead_attn_out_proj_bias = self.layers_0_multihead_attn_out_proj_bias
  layers_0_multihead_attn_out_proj_weight = self.layers_0_multihead_attn_out_proj_weight
  layers_0_multihead_attn_in_proj_bias = self.layers_0_multihead_attn_in_proj_bias
  layers_0_multihead_attn_in_proj_weight = self.layers_0_multihead_attn_in_proj_weight
  layers_0_norm1_bias = self.layers_0_norm1_bias
  layers_0_norm1_weight = self.layers_0_norm1_weight
  layers_0_self_attn_out_proj_bias = self.layers_0_self_attn_out_proj_bias
  layers_0_self_attn_out_proj_weight = self.layers_0_self_attn_out_proj_weight
  layers_0_self_attn_in_proj_bias = self.layers_0_self_attn_in_proj_bias
  layers_0_self_attn_in_proj_weight = self.layers_0_self_attn_in_proj_weight
  _0 = ops.profiler._record_function_enter("FlattenModule")
  _1 = torch.transpose(tensor_0, 1, 0)
  _2 = torch.t(layers_0_self_attn_in_proj_weight)
  _3 = torch.addmm(layers_0_self_attn_in_proj_bias, torch.view(_1, [10, 512]), _2)
  _4 = torch.split(torch.view(_3, [10, 1, 1536]), 512, -1)
  _5, _6, _7, = _4
  _8 = torch.view(torch.clone(_5, memory_format=0), [10, 8, 64])
  _9 = torch.transpose(_8, 0, 1)
  _10 = torch.view(torch.clone(_6, memory_format=0), [10, 8, 64])
  _11 = torch.transpose(_10, 0, 1)
  _12 = torch.view(torch.clone(_7, memory_format=0), [10, 8, 64])
  _13 = torch.transpose(_12, 0, 1)
  _14 = torch.bmm(torch.div(_9, CONSTANTS.c0), torch.transpose(_11, -2, -1))
  _15 = torch.bmm(torch._softmax(_14, -1, False), _13)
  _16 = torch.clone(torch.transpose(_15, 0, 1), memory_format=0)
  _17 = torch.view(_16, [10, 512])
  _18 = torch.t(layers_0_self_attn_out_proj_weight)
  _19 = torch.addmm(layers_0_self_attn_out_proj_bias, _17, _18)
  _20 = torch.transpose(torch.view(_19, [10, 1, 512]), 1, 0)
  _21, _22, _23 = torch.native_layer_norm(torch.add(tensor_0, _20), [512], layers_0_norm1_weight, layers_0_norm1_bias, 1.0000000000000001e-05)
  _24 = torch.transpose(_21, 1, 0)
  _25 = torch.transpose(tensor_1, 1, 0)
  _26 = torch.split_with_sizes(layers_0_multihead_attn_in_proj_weight, [512, 1024])
  _27, _28, = _26
  _29 = torch.split_with_sizes(layers_0_multihead_attn_in_proj_bias, [512, 1024])
  _30, _31, = _29
  _32 = torch.t(_27)
  _33 = torch.addmm(_30, torch.view(_24, [10, 512]), _32)
  _34 = torch.view(_33, [10, 1, 512])
  _35 = torch.t(_28)
  _36 = torch.addmm(_31, torch.view(_25, [20, 512]), _35)
  _37 = torch.split(torch.view(_36, [20, 1, 1024]), 512, -1)
  _38, _39, = _37
  _40 = torch.transpose(torch.view(_34, [10, 8, 64]), 0, 1)
  _41 = torch.view(torch.clone(_38, memory_format=0), [20, 8, 64])
  _42 = torch.transpose(_41, 0, 1)
  _43 = torch.view(torch.clone(_39, memory_format=0), [20, 8, 64])
  _44 = torch.transpose(_43, 0, 1)
  _45 = torch.bmm(torch.div(_40, CONSTANTS.c0), torch.transpose(_42, -2, -1))
  _46 = torch.bmm(torch._softmax(_45, -1, False), _44)
  _47 = torch.clone(torch.transpose(_46, 0, 1), memory_format=0)
  _48 = torch.view(_47, [10, 512])
  _49 = torch.t(layers_0_multihead_attn_out_proj_weight)
  _50 = torch.addmm(layers_0_multihead_attn_out_proj_bias, _48, _49)
  _51 = torch.transpose(torch.view(_50, [10, 1, 512]), 1, 0)
  _52, _53, _54 = torch.native_layer_norm(torch.add(_21, _51), [512], layers_0_norm2_weight, layers_0_norm2_bias, 1.0000000000000001e-05)
  _55 = torch.t(layers_0_linear1_weight)
  _56 = torch.addmm(layers_0_linear1_bias, torch.view(_52, [10, 512]), _55)
  _57 = torch.relu(torch.view(_56, [1, 10, 2048]))
  _58 = torch.t(layers_0_linear2_weight)
  _59 = torch.addmm(layers_0_linear2_bias, torch.view(_57, [10, 2048]), _58)
  _60 = torch.add(_52, torch.view(_59, [1, 10, 512]))
  _61, _62, _63 = torch.native_layer_norm(_60, [512], layers_0_norm3_weight, layers_0_norm3_bias, 1.0000000000000001e-05)
  ops.profiler._record_function_exit(_0)
  return _61
```

I know that the above error may ultimately affect the accuracy of the model, so I want to confirmed: what could be the specific reasons for the above error?

### Versions

version:
Torch 1.13.0
cuda11.7

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel