Untitled

        idx=torch.searchsorted(torch.softmax(-150 * probs, 0).cumsum(0), torch.rand(500))
        idx=idx[torch.where(idx<self.so3.shape[0])]
        fake_Rs[0:len(idx)] = self.so3[idx]
        idx=torch.searchsorted(torch.softmax(-15 * probs, 0).cumsum(0), torch.rand(1500))
        idx=idx[torch.where(idx<self.so3.shape[0])]
        fake_Rs[500:500+len(idx)]=self.so3[idx]
        idx=torch.searchsorted(torch.softmax(-5 * probs, 0).cumsum(0), torch.rand(1000))
        idx=idx[torch.where(idx<self.so3.shape[0])]
        fake_Rs[2000:2000+len(idx)]=self.so3[idx]