hello all I’d like to write the function below in pytorch in order to fully utilize my computer’s gpu such that I can speed up the calculation. Currently this step is really slowing my algorithm down, any help is greatly appreciated. Some notes, the function essentially creates a gaussian kernel with a predefined window size and sigma which is created with respect to the euclidean distance between the center of the input feature map and the current point of interest.
def pseudo_colliculus_map (neural_map):
    dimensions = neural_map.size()
    mapped = torch.zeros(dimensions[2], dimensions[3]).cuda()
    localizer = torch.Tensor([(dimensions[2]/2),(dimensions[3]/2)]).type(torch.int64).cuda()
    for i in range(0, dimensions[2]):
       for j in range(0, dimensions[3]):

            euc = math.sqrt((localizer[0]-i)**2 + (localizer[1]-j)**2)
            sigma = 0.06*(euc) + 0.4
            size = torch.Tensor(np.array(sigma*3)).type(torch.int64).cuda()
            kernel = np.fromfunction(lambda x, y: (1/(2*math.pi*sigma**2)) * math.e ** ((-1*((x-(size-                         1)/2)**2+(y-(size-1)/2)**2))/(2*sigma**2)), (size, size))
           kernel = torch.div( kernel, torch.sum(kernel))
        
        
            weights = kernel.type(torch.FloatTensor).cuda()
            weights = weights.view(1,1,size, size)
            pad = math.ceil((size)/2)
            mapping = F.conv2d(neural_map, weights, stride =1, padding = pad)
            val = mapping[0, 0, i,j]
        
            mapped[i,j] = val



    return mapped