In your forward_ssl code, you squeeze the fps sampling sequence, but this will raise errors when the batch size is 1 because you index the fps_choice with the batch index afterward.
fps_choice = pointnet2_batch_utils.furthest_point_sample(
points_overlap_org.contiguous(), self.npoints
).long().squeeze()
# compute correspondence
for kk in range(batch_size):
key_idx = example['overlap'][kk][fps_choice[kk]].type(torch.float32)
example['correspondence'][kk] = torch.min(torch.abs(torch.sub(key_idx.unsqueeze(dim=-1),example['correspondence'][kk].unsqueeze(dim=0).type(torch.float32))),dim=-1)[1]