Python - многопроцессорность застревает

Я пытаюсь параллельно выполнять определенные операции с набором данных MNIST. С многопроцессорностью он застревает и никогда не заканчивается. Я попытался поместить функции закрытия и соединения внутри или вне цикла for. Я также попытался поместить экземпляр пула внутри функции, похоже, тоже не работает. Что я здесь делаю не так? Не знаю, где мне искать, чтобы решить эту проблему. Большое спасибо за вашу помощь заранее

class MnistPartialDataset(Dataset):
    def __init__(self, data, transform=None):
        self.data = data
    def __getitem__(self, index):
    return self.data[index]
    def __len__(self):
        return len(self.data)

def create_tasks(train_pool, test_pool):
    train_thread_result = []
    start = time.time()
    for k, v in train_tasks.items():
        tc_data = train_pool.map_async(test_function, v)
        train_thread_result.append(tc_data)
    train_pool.close()
    train_pool.join()
    for i in range(len(train_thread_result)):
        train_dataset[i] = MnistPartialDataset(train_thread_result[i].get())
    end = time.time()
    print("tc time: " + str(end-start))

    test_thread_result = []
    start = time.time()
    for k, v in test_tasks.items():
        tc_data = test_pool.map_async(test_function, v)
        test_thread_result.append(tc_data)
    test_pool.close()
    test_pool.join()
    for i in range(len(test_thread_result)):
        test_dataset[i] = MnistPartialDataset(test_thread_result[i].get())
    end = time.time()
def main():
    train_pool =  multiprocessing.Pool(processes = multiprocessing.cpu_count()-1)
    test_pool =  multiprocessing.Pool(processes = multiprocessing.cpu_count()-1)
    create_tasks(train_pool, test_pool)
    pool.close()
    pool.terminate()