Solved by Vasiliy_Kuznetsov in post #2 
              
              
                Hi @keyky, this is a limitation of symbolic tracing with FX.  Here is a workaround using torch.fx.wrap: 
@torch.fx.wrap                                                           
def shift_left(out, x, fold):                                            
    out[:, :-1, :fold] = x[:, 1:, :fold]  # shi…