This piece of code can do the job. But it’s not very optimal.
import torch

n = 4
a = 1

A = torch.zeros((n, n))

for i in range(n) :
    for j in range(max(i - a, 0), min(i + a + 1, n)) :
        A[i][j] = 1.0