Fix mask issue in distributed nt_xent loss

[lxysl]

Fix the mask issue #30 in nt_xent loss when training in a distributed situation. The all_gather result of z used to build mask should be something like [z_i1, z_i2, ..., z_iw, z_j1, z_j2, ..., z_jw] instead of [z_i1, z_j1,z_i2, z_j2, ..., z_iw, z_jw], where w represents the world_size.