numpy

def gen(out,griding_num = 100,localization_type='rel'):
out = out.data.cpu().numpy()
out_loc = np.argmax(out,axis=0)

if localization_type == 'rel':

    out =out[:-1, :, :]
    out = np.exp(out) / np.sum(np.exp(out), axis=0)
    prob = out
    print(prob)


    idx = np.arange(griding_num)
    idx = idx.reshape(-1, 1, 1)

    loc = np.sum(prob * idx, axis=0)

    loc[out_loc == griding_num] = griding_num
    out_loc = loc
lanes = []
for i in range(out_loc.shape[1]):
    out_i = out_loc[:,i]
    lane = [int(round((loc + 0.5) * 1280.0 / (griding_num - 1))) if loc != griding_num else -2 for loc in out_i]
    lanes.append(lane)
return lanes