There you go, this will be much faster

Note that I use `.narrow(-1, 0, x_size-1)`

out of habit but `[:, :-1]`

works as well if you prefer that notation.

```
import torch
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
print(A)
A = torch.zeros((3, 3, 3), dtype = torch.float)
# This code assumes A is contiguous ! If it is not, add
# A = A.contiguous()
# For indexing, collapse the last two dimensions of A
A_view = A.view(A.size(0), -1)
# Compute the indices where you will index in A
x_size = X.size(-1)
indices = X.narrow(-1, 0, x_size-1) * A.stride(1) * A.stride(2) + X.narrow(-1, 1, x_size-1) * A.stride(2)
# Put 1s at the computed indices
A_view.scatter_(1, indices, 1)
print(A)
```