Why the torch_load_kwargs parameter doesn't work?

[githubnemo]

Discussed in #1090

^{Originally posted by lnzpgo January 26, 2025}
I'm getting these annoying warning messages from torch.load:

/home/user/Code/ML-PROJ/.venv/lib/python3.11/site-packages/skorch/net.py:2261: FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for weights_only will be flipped to True. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via torch.serialization.add_safe_globals. We recommend you start setting weights_only=True for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.

Despite I'm explicitely passing weights_only=True to skorch...

# Wrap the PyTorch model in a scikit-learn compatible interface using skorch
net = NeuralNetRegressor(
    module=SimpleNN,
    criterion=nn.MSELoss,
    optimizer=optim.Adam,
    max_epochs=20,
    batch_size=32,
    train_split=None,  # Disable internal validation split (we use cross-validation)
    verbose=0,
    # Set weights_only=True for loading model weights
    torch_load_kwargs={'weights_only': True},
)
# Define hyperparameter grids for Grid Search and Random Search
param_grid = {
    'module__hidden_size': [10, 20, 30],  # Number of units in the hidden layer
    'optimizer__lr': [0.001, 0.01, 0.1],  # Learning rate
    'batch_size': [32, 64, 128, 256],  # Batch size
}
# Grid Search with Cross-Validation
print("Performing Grid Search...")
grid_search = GridSearchCV(
    estimator=net,
    param_grid=param_grid,
    scoring='neg_mean_squared_error',  # Use negative MSE for scoring
    cv=5,  # 5-fold cross-validation
    n_jobs=-1,  # Use all available CPU cores
    verbose=1,
)
grid_search.fit(X, y)

Reason

NeuralNet.__setstate__ currently ignores self.torch_load_kwargs but shouldn't.