Save memory when compute shortest distance about alignedreid HOT 3 OPEN

def shortest_path(dist_mat):
    """
    Args:
    dist_mat: pytorch Variable, available shape:
      1) [m, n]
      2) [m, n, N], N is batch size
      3) [m, n, *], * can be arbitrary additional dimensions
      Returns:
    dist: three cases corresponding to `dist_mat`:
      1) scalar
      2) pytorch Tensor, with shape [N]
      3) pytorch Tensor, with shape [*]
    """
    if dist_mat.dim() == 3:
        M, N, B = dist_mat.size()
        helper = torch.zeros(N, B)
    elif dist_mat.dim() == 2:
        M, N = dist.size()
        helper = torch.zeros(N)
    for i in range(M):
        for j in range(N):
            if i == 0 and j == 0:
                helper[j] = dist_mat[i][j]
            elif i == 0:
                helper[j] = helper[j-1] + dist_mat[i][j]
            elif j == 0:
                helper[j] = helper[j] + dist_mat[i][j]
            else:
                helper[j] = torch.min(helper[j-1], helper[j]) + dist_mat[i][j]
    return helper[-1]