Why the torch_load_kwargs parameter doesn't work?

[lnzpgo]

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)

[githubnemo]

Thanks for the question. I can see how this is going to be annoying quickly.

The reason this is happening is because __setstate__ currently ignores torch_load_kwargs. I think there's no good reason why we shouldn't apply those kwargs here as well.

[BenjaminBossan]

Thanks for the report @lnzpgo. #1092 should resolve the issue. If you install skorch from source, the warning should go away. However, this requires to load the net and save it again. The first load would still show the warning, but then it should be gone.

Alternatively, you may want to save and load only the params, which honors the torch_load_kwargs.

[lnzpgo]

Thank you for your work.

In the meantime, I have found this short term solution (to be put at the beginning of the file):

import os
os.environ["PYTHONWARNINGS"] = "ignore:You are using `torch.load` with `weights_only=False`"