Thanks for your answer.
Let’s note a function named g. g takes as input two tensors : centroid which is of size (n,2) and grid which is of size (m,2)
g will return a size (m) tensor called ‘res’ for which each value of ‘res’ is the distance of the nearest point for each point in the grid to all centroids

And that is done inside the for loop.
The for loop is there because ideally I’d like to take multiple grids and multiple centroids and use element wise operation to make the computation quick. There number of multiple grids and multiple centroids is denoted by t.
In your example, centroids would be
centroids = torch.rand(t,n,2)
But the problem is that n may vary for each t, so I think centroids must be a list of tensors, or find a way to use a mask somehow so that the function knows how many centroids it has to deal with.
Sorry if it’s still confusing…