How to make a 2D tensor whose diagonal and its neighbors are ones

1

I want to make a 2D tensor whose diagonal and its neighbors are filled with ones and others are zeros.

For example 5x5 matrix and neighbor 1,
[[1, 1, 0, 0, 0],
[1, 1, 1, 0, 0],
[0, 1, 1, 1, 0],
[0, 0, 1, 1, 1],
[0, 0, 0, 1, 1]]

Is there any pytorch operation for making those matrix?

2 
n = 5
torch.diagflat(torch.ones(n-1), 1) + torch.diagflat(torch.ones(n-1), -1) + torch.eye(n)
1 Like