I have written pytorch custom loss, but it ran too slow. Here is the code I wrote.

def exp_sum(refer_h, others, sigma):
    out = list()

    for item in others:
        tmp = torch.exp(- (refer_h - item) ** 2 / (2 * sigma ** 2))
        tmp = torch.mean(tmp)
        out.append(tmp)

    return sum(out)


Any advice to speed up my program? Thanks