Pytorch已知一个tensor的索引，并对它进行重新排序

之前遇到了一个问题，我有一个tensor,它的形状是b c h w ，和与它一样长的一个索引 b c 1 1 。我们要根据这个索引对原始tensor进行重新排序。百度了一圈貌似没有直接解决的api,感觉只能用循环处理，因此我用循环写了一个函数来搞定这个问题。
#第一个参数是原始要排序的tensor,和索引，返回排好序的结果。

def sort_with_idx(source,idx):
    #根据索引对tensor进行排序
    "source: b c h w ,  idx: b c 1 1"
    b,c,_,_ = source.size()
    after_sort = torch.zeros_like(source)
    idx = idx[:,:]
    for i in range(b):
        for j in range(c):
            after_sort[i][j][:][:] = source[i][idx[i][j]][:][:]
    return after_sort

我已经对他进行了测试，是能够正常按照要求排序的。

a和b的值