Attribute Error when using GridSearchCV and Pytorch + SklearnPipelines

I defined a class TorchRegressionModule using pytorch and then later I created a pipe line and tried to pass it through GridSearchCV as estimator but I got the following error there.

f16_elevators_pipeline_steps = [
    ("selector", sklearn.feature_selection.SelectKBest(
        score_func=sklearn.feature_selection.mutual_info_regression,
        k=desired_features
    )),
    ("scaler", sklearn.preprocessing.MinMaxScaler()),
    ("estimator", skorch.NeuralNetRegressor(
        TorchRegressionModule,
        optimizer=torch.optim.Adam,
        lr=0.01,
        max_epochs=50,
        callbacks=[skorch.callbacks.EarlyStopping(
            patience=10)],
        batch_size=10,
    ))
]

Then if I try to pass it under GridSearchCV,

> grid = GridSearchCV(f16_elevators_pipeline, param_grid=param_grid, n_jobs=-1, cv=3)
> grid_result = grid.fit(f16_elevators_X_train.to_numpy().astype('float32'), f16_elevators_y_train.values.astype('float32').reshape(-1,1))
> AttributeError: Can't get attribute 'TorchRegressionModule' on <module '__main__' (built-in)>
> """
> 
> 
> The above exception was the direct cause of the following exception:

Can anyone let me know what should I do to fix this issue, as I am unsure of what to do here, so that I can pass my pipeline with Neural network as one of the arguments in pipeleine?