Please help me to vectorize this code on the batch dimension at least:
A = torch.zeros((3, 3, 3), dtype = torch.float)
X = torch.tensor([[0, 1, 2, 0, 1, 0], [1, 0, 0, 2, 1, 1], [0, 0, 2, 2, 1, 1]])
for a, x in zip(A, X):
    for i, j in zip(x, x[1:]):
        a[i, j] = 1

Thanks!