Convert TSM pytorch model to onnx

I have succees converting the tsm pytorch to onnx. but the onnx output is different with the pytorch. when converting. it show me that:

TracerWarning: Converting a tensor to a Python index 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!
  out[:, :-1, :fold] = x[:, 1:, :fold]  # shift left
/Users/zhanghongxing/vscode_project/tsm_inference/infer/../ops/temporal_shift.py:37: TracerWarning: There are 2 live references to the data region being modified when tracing in-place operator copy_ (possibly due to an assignment). This might cause the trace to be incorrect, because all other views that also reference this data will not reflect this change in the trace! On the other hand, if all other views use the same memory chunk, but are disjoint (e.g. are outputs of torch.split), this might still be safe.
  out[:, :-1, :fold] = x[:, 1:, :fold]  # shift left
/Users/zhanghongxing/vscode_project/tsm_inference/infer/../ops/temporal_shift.py:38: TracerWarning: Converting a tensor to a Python index 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!
  out[:, 1:, fold: 2 * fold] = x[:, :-1, fold: 2 * fold]  # shift right
/Users/zhanghongxing/vscode_project/tsm_inference/infer/../ops/temporal_shift.py:38: TracerWarning: There are 2 live references to the data region being modified when tracing in-place operator copy_ (possibly due to an assignment). This might cause the trace to be incorrect, because all other views that also reference this data will not reflect this change in the trace! On the other hand, if all other views use the same memory chunk, but are disjoint (e.g. are outputs of torch.split), this might still be safe.
  out[:, 1:, fold: 2 * fold] = x[:, :-1, fold: 2 * fold]  # shift right
/Users/zhanghongxing/vscode_project/tsm_inference/infer/../ops/temporal_shift.py:39: TracerWarning: Converting a tensor to a Python index 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!
  out[:, :, 2 * fold:] = x[:, :, 2 * fold:]  # not shift
/Users/zhanghongxing/vscode_project/tsm_inference/infer/../ops/temporal_shift.py:39: TracerWarning: There are 2 live references to the data region being modified when tracing in-place operator copy_ (possibly due to an assignment). This might cause the trace to be incorrect, because all other views that also reference this data will not reflect this change in the trace! On the other hand, if all other views use the same memory chunk, but are disjoint (e.g. are outputs of torch.split), this might still be safe.
**I didnot use the inplace shift**
  out[:, :, 2 * fold:] = x[:, :, 2 * fold:]  # not shift

        out[:, :-1, :fold] = x[:, 1:, :fold] # shift left
        out[:, 1:, fold: 2 * fold] = x[:, :-1, fold: 2 * fold]  # shift right
        out[:, :, 2 * fold:] = x[:, :, 2 * fold:] # not shift

I think this may be caused by the shift operation. but I dont think it is a inplace op. it seems that the onnx not support the slice op.

[Usernamezhx]

@ChiShao you can change https://github.com/mit-han-lab/temporal-shift-module/blob/f09f42db80f1dbeaf9c7448fbd491cd59043e711/ops/temporal_shift.py#L27 to this:

    def shift(x, n_segment, fold_div=3):
        nt, c, h, w = x.size()
        n_batch = nt // n_segment
        x = x.view(n_batch, n_segment, c, h, w)

        fold = c // fold_div

        left_side = torch.cat((x[:, 1:, :fold], torch.zeros(1, 1, fold, h, w)), dim=1)
        middle_side = torch.cat((torch.zeros(1, 1, fold, h, w), x[:, :n_segment - 1, fold: 2 * fold]), dim=1)
        out = torch.cat((left_side, middle_side, x[:, :, 2 * fold:]), dim=2)

        return out.view(nt, c, h, w)

[bujianyiwang]

@Usernamezhx can you provide the code of coverting to onnx?