How to right shift each row of a matrix?

autograd
1

I have a matrix whose shape is (TxK, and K << T). I want to extend it into shape TxT, and right shift the i-th row with i steps.

For an example:

inputs: T= 5, and K = 3
1 2 3
1 2 3
1 2 3
1 2 3
1 2 3

expected outputs:
1 2 3 0 0 0
0 1 2 3 0 0
0 0 1 2 3 0
0 0 0 1 2 3
0 0 0 0 1 2

My solutions:

right_pad = T - K + 1
output = F.pad(input, (0, right_pad), 'constant', value=0)
output = output.view(-1)[:-T].view(T, T)

My solution will cause the error – gradient computation has been modified by an inplace operation. Is there an efficient and feasible way to achieve my purpose?