如何在我的代码中并行处理此瓶颈？

如何解决如何在我的代码中并行处理此瓶颈？

我正在使用pytorch进行数据处理，而不是用于机器学习等。这大大提高了性能，但是我仍然遇到瓶颈。大多数时候，我在大型数组上执行简单的操作，包括布尔索引。但是，在某些时候，我需要将“最终”数组与另一个数据集中的值进行比较，并找到匹配项。像这样：

import numpy as np
import torch 

device = torch.device("cuda:0")

# make couple tensors
data1 = torch.rand((1000,10),device=device)
data2 = torch.rand((1000,device=device)

# do some arithmetic operations on tensors (real code has a lot more operations)
data3 = data1 * data2

# a dataframe was built previously and must be loaded.  Here a dummy dataframe is created
df = pd.DataFrame(data=np.random.randint(0,10,(10,5)),columns=list("ABCDE"))

# define flag array
matchFound = np.full((8000,1),False)


# now the bottleneck
for row,val in df.A.unique():
    # determine if anything in data3 is within 10% of "val"
    idx = (val >= 0.9 * data3) & (val <= 1.1 * data3)

    # if there are any matches then set flag value
    if idx.sum() > 0:
        matchFound[row,0] = True
    
    

所有达到瓶颈点的过程都需要3-4毫秒，但是瓶颈需要750毫秒。我需要经历〜18E6次，使总运行时间为几周。有什么办法可以加快速度吗？

到目前为止，我已经尝试了什么：我编写了一个函数来封装瓶颈区域，并使用来自多处理的Process将其传递给多个CPU（我也尝试使用来自torch.multiprocessing的Process）。想法是将该数据帧分为15个块，然后将每个块发送到CPU，从而将瓶颈时间从750 ms减少到50 ms，并将我的总运行时间减少到1d。但是，我一直收到运行时错误（RuntimeError：CUDA错误：初始化错误）。我不明白是什么导致了此错误（将张量传递给多个CPU时CUDA会生气吗？）。另外，我认为通过将数据从GPU发送到CPU所花费的时间比节省的时间还要多。