The new_group API requires all processes to call with the same ranks argument if even if they do not participate in the new group. See the API doc here: https://pytorch.org/docs/stable/distributed.html#torch.distributed.new_group 8
In the code above, the following code breaks the above assumption.
    if rank == 0 or rank == 1:
        ranks = [0,1]
    if rank == 2 or rank == 3:
        ranks = [2,3]
    print('ranks: ', ranks)
    group = dist.new_group(ranks = ranks)

It needs to be modified to the following:
    g1 = dist.new_group(ranks = [0, 1])
    g2 = dist.new_group(ranks = [2, 3])
    # check rank to see if the current process should use g1 or g2