import torch
A=torch.ones(2,3,4,2).long()
B=torch.ones(2,3,4,2).long()    

N,H,W,_ = A.shape
b = torch.arange(N).view(N,1,1) #N,1,1
j = torch.arange(H)
k = torch.arange(W)
jk = torch.meshgrid(j,k)
jk = torch.stack(jk,2).view(1,H,W,2) #1,H,W,2
A[b,B[...,0],B[...,1]] = jk # broadcast to N,H,W,2