n_jobs> 1 использование sklearn и pytorch возможно внутри Neuraxle?

Я построил свой собственный склеарноподобный оценщик, используя обучение pytorch внутри графического процессора (cuda), и он отлично работает с RandomizedSearchCV, когда n_jobs == 1. Когда n_jobs> 1, я получаю следующую ошибку:

PicklingError: не удается pickle main .LSTM'>: поиск атрибута LSTM на main не удался

Это фрагмент кода, который дает мне ошибку:

model = my_model(input_size=1, hidden_layer_size=80, n_lstm_units=3, bidirectional=False,
                 output_size=1, training_batch_size=60, epochs=7500, device=device)
model.to(device)

hidden_layer_size = random.uniform(40, 200, 20).astype("int")
n_lstm_units = arange(1, 4)

parametros = {'hidden_layer_size': hidden_layer_size, 'n_lstm_units': n_lstm_units}

splitter = ShuffleSplit()

regressor = model
cv_search = \
    RandomizedSearchCV(estimator=regressor, cv=splitter,
                  search_spaces=parametros,
                  refit=True,
                  n_iter=4,
                  verbose=1,
                  n_jobs=2,
                  scoring=make_scorer(mean_squared_error,
                                      greater_is_better=False,
                                      needs_proba=False))

cv_search = MetaSKLearnWrapper(cv_search)
cv_search.fit(X, y)

Использование оболочки Neuraxle приводит к точно такой же ошибке. , ничего не меняет.

Я нашел ближайшее решение здесь, но до сих пор не знаю, как использовать RandomizedSearchCV в Neuraxle . Это совершенно новый проект, поэтому я не смог найти ответа ни в их документации, ни в примерах сообщества. Если кто-нибудь может привести мне пример или хорошее указание, это спасет мне жизнь. Спасибо

Ps: Также помогает любой способ запустить RandomizedSearchCV с моей моделью pytorch на графическом процессоре без Neuraxle, мне просто нужно n_jobs> 1.

Ps2: В моей модели есть метод fit (), который создает и перемещает тензоры в графический процессор и уже протестирован.