Find neighborhood for torch tensor

1

I am trying to implement a Self-Organizing Map where for a given input sample, the best matching unit/winning unit is chosen based on (say) L2-norm distance between the SOM and the input. The winning unit/BMU (som[x, y]) has the smallest L2 distance from the given input (z):

# Input batch: batch-size = 512, input-dim = 84-
z = torch.randn(512, 84)

# SOM shape: (height, width, input-dim)-
som = torch.randn(40, 40, 84)

print(f"BMU row, col shapes; row = {row.shape} & col = {col.shape}")
# BMU row, col shapes; row = torch.Size([512]) & col = torch.Size([512])

For clarity, for the first input sample in the batch “z[0]”, the winning unit is “som[row[0], col[0]]”-

z[0].shape, som[row[0], col[0]].shape
# (torch.Size([84]), torch.Size([84]))

torch.norm((z[0] - som[row[0], col[0]])) is the smallest L2 distance between z[0] and all other som units except row[0] and col[0].

# Define initial neighborhood radius and learning rate-
neighb_rad = torch.tensor(2.0)
lr = 0.5

# To update weights for the first input "z[0]" and its corresponding BMU "som[row[0], col[0]]"-
for r in range(som.shape[0]):
    for c in range(som.shape[1]):
        neigh_dist = torch.exp(-torch.norm(input = (som[r, c] - som[row[0], col[0]])) / (2.0 * torch.pow(neighb_rad, 2)))
        som[r, c] = som[r, c] + (lr * neigh_dist * (z[0] - som[r, c]))

How can I implement the code for:

  1. updating weights for all units around each BMU without the 2 for loops (and)
  2. do it for all of the inputs “z” (here, z has 512 samples)