914192dda940237aaf47130e5511bb929cd776b7,torchsample/utils.py,,th_gather_nd,#Any#Any#,252

Before Change

    if coords.size(1) != x.dim():
        raise ValueError("Coords must have column for each image dim")

    inds = coords[:,0]*x.size(1)
    for i in range(x.dim()-2):
        inds += coords[:,i+1]*x.size(i+2)
    inds += coords[:,-1]
    x_gather = torch.index_select(th_flatten(x), 0, inds)
    return x_gather

After Change

def th_gather_nd(x, coords):
    inds = coords.mv(torch.LongTensor(x.stride()))
    x_gather = torch.index_select(th_flatten(x), 0, inds)
    return x_gather

Instances