This sounds like a data parallel problem, although the nested structure seems like it might be an issue if you need to “unpack” PyTorch tensors from Python data structures.
For this case, you might want to consider preprocessing the data into a higher dimensional tensor, e.g., 100000,2,300.
From here, it becomes a straightforward data-parallel problem that shouldn’t take long on 4 GPUs (or likely even a single GPU) if you move the split data into equal sized parts, move the different parts to different GPUs, and aggregate the results.
At this point I’d be concerned that all of the preprocessing time (e.g., the steps to generate tmp_df and pack it into a large tensor) would be significantly greater than the calculation time for cosine similarity.