@gphilip Thanks for your response. I’ve been going off the example detailed here, but I have a few questions about testing this out. Here is what I have so far:
from torch import nn
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


class CosSimNetwork(nn.Module):
    def __init__(self):
        super(CosSimNetwork, self).__init__()

    def forward(self, embeds):
        #Here do the cosine similarity calculation
        cos_sims = 12
        return cos_sims

def calc_cos_sims(rank, world_size):
    dist.init_process_group('gloo', rank=rank, world_size=world_size)
    model = CosSimNetwork()
    ddp_model = DDP(model, device_ids=[rank])

    cos_sims = ddp_model(DATA.to(rank))


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()

If I understand correctly, because I have 4 GPUs world_size = 4. However, I’m not entirely sure what rank should be given that I have 4 GPUs. Also, I have my data stored in a pickled pandas dataframe with two columns being the embeddings that would be used in the cosine similarity calculation. It looks something like cos_sim(embed_a, embed_b). I’m not entirely sure how to read in the data properly for DistributedDataParallel and the forward method of CosSimNetwork. Any advice is much appreciated.