I’m wanting to use pytorch for it’s tensor math and not necessarily for training a ML model. I have a machine that has 4 GPUs on them. The overall goal is to calculate the cosine similarity between pairs of embeddings utilizing all 4 GPU. Here is a walkthrough of what i have thus far:
Read the data in as a pandas dataframe
import pandas as pd
tmp_df = pd.read_pickle('./embed_pairs_df.pkl')

There are 2 columns to tmp_df: embeds_a, and embeds_b. Each column contains 1 million rows making tmp_df.shape = (100000,2). Each element in a row is a pytorch tensor that has shape (300,1). I’ve already worked out the combinations of pairs and stored them appropriately in tmp_df. As such, tmp_df['embeds_a'][i] and tmp_df['embeds_b'][i] would constitute the ith pair of embeddings that would need to be run through a cosine similarity function.
My question is, given the format of the data and the 4 GPUs available, what is the best way to distribute the cosine similarity calculation across all 4 GPUs in parallel?