The following code works:
import torch

# randomly produces a 1-D permutation index array,
# such that each element of the shuffled array has
# a distance less than K from its original location
def rand_perm(n, K):
    o = torch.arange(0, n)
    if K <= 1:
        return o
    while True:
        p = torch.randperm(n)
        d = abs(p - o) < K
        if bool(d.all()):
            return p

if __name__ == '__main__':
    for i in range(10):
        print(rand_perm(10, 2))

but it seems that when n is large, and K is small, the generation will take a very long time. Does there exists a more efficient implementation?