Torch.jit.trace get the wrong model

[Dream-math]

Does anybody try to convert this model to ScriptModule by torch.jit.trace, we finally get it, but given the same input, the generated script module get the wrong result. During converting there are some messages like below:

/home/hermanhe/.local/lib/python3.10/site-packages/torch/functional.py:478: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2894.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/restorers/srgan.py:94: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if test_mode:
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/real_basicvsr_net.py:83: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if torch.mean(torch.abs(residues)) < self.dynamic_refine_thres:
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:118: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert h >= 64 and w >= 64, (
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:73: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if lrs.size(1) % 2 == 0:
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:334: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  w_up = w if (w % 32) == 0 else 32 * (w // 32 + 1)
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:335: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  h_up = h if (h % 32) == 0 else 32 * (h // 32 + 1)
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:335: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
  h_up = h if (h % 32) == 0 else 32 * (h // 32 + 1)
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:296: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
  flow = ref[0].new_zeros(n, 2, h // 32, w // 32)
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/common/flow_warp.py:27: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if x.size()[-2:] != flow.size()[1:3]:
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/common/flow_warp.py:41: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  grid_flow_x = 2.0 * grid_flow[:, :, :, 0] / max(w - 1, 1) - 1.0
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/common/flow_warp.py:42: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  grid_flow_y = 2.0 * grid_flow[:, :, :, 1] / max(h - 1, 1) - 1.0
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:352: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  flow[:, 0, :, :] *= float(w) / float(w_up)
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:353: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  flow[:, 1, :, :] *= float(h) / float(h_up)
/home/hermanhe/.local/lib/python3.10/site-packages/mmedit/models/backbones/sr_backbones/basicvsr_net.py:132: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if i < t - 1:  # no warping required for the last timestep

Any advice about getting the right script module? Thanks!

[Feynman1999]

Are you using torch.jit.trace or torch.jit.script?