@pritamdamania87 Here is what I now have:
import torch
import torch.multiprocessing as mp
import torch.distributed as dist
import torch.nn.functional as F
import pandas as pd


def calc_cos_sims(rank, world_size):
    dist.init_process_group('gloo', rank=rank, world_size=world_size)
    cuda_device = torch.device('cuda:'+str(rank))
    data_path = './embed_pairs_df_million_part_' + str(rank) + '.pkl'
    tmp_df = pd.read_pickle(data_path)

    embeds_a_list = [embed_a for embed_a in tmp_df['embeds_a']]
    embeds_b_list = [embed_b for embed_b in tmp_df['embeds_b']]

    embeds_a_tensor = torch.tensor(embeds_a_list, device=cuda_device)
    embeds_b_tensor = torch.tensor(embeds_b_list, device=cuda_device)
    cosine_tensor = F.cosine_similarity(embeds_a_tensor, embeds_b_tensor)
    


def main():
    world_size = 4 #since I have 4 GPUs on a single machine
    mp.spawn(calc_cos_sims,
             args=(world_size,),
             nprocs=world_size,
             join=True)

if __name__ == 'main':
    main()

This has changed to assume that the data has been split up into 4 different parts (1 part for each GPU) by data_path = './embed_pairs_df_million_part_' + str(rank) + '.pkl'. Since each entry in a column is a numpy array, I went ahead and converted everything to pytorch tensors. cosine_tensor is the cosine similarity between each element of the data split. I read the link you posted about aggregating, but I’m not entirely sure how to implement it. How would that be done in this case?
Also, have I called mp.spawn correctly? I have args=(world_size,) but calc_cos_sims(rank, world_size) has the rank first followed by the world size and those go into dist.init_process_group('gloo', rank=rank, world_size=world_size).