Repository: pytorch
Issue: Inconsistent Model Loading with `weights_only=True` for Custom Models
Labels: high priority,needs reproduction,module: serialization,triaged

### 🚀 The feature, motivation and pitch

I'm encountering an issue with the `weights_only=True` parameter in `torch.load`, where I need to load custom model classes while ensuring repeatable results across multiple runs. My current workaround involves setting `weights_only=False` to achieve identical predictions, which raises security warnings due to potential arbitrary code execution. This issue serves to inform the PyTorch team about the challenges in achieving both secure and repeatable model loading in use cases like neural architecture search, where multiple models are saved and loaded regularly.

**Background**:

In my work with neural architecture search, achieving repeatability is essential. When loading a saved custom model using `weights_only=True`, I encounter issues with loading state and repeatability, even after adding custom classes to `torch.serialization.add_safe_globals`. Using `weights_only=False` resolves the issue, but it raises security concerns. Given the importance of security and repeatability, this feature request aims to address these limitations and propose an improvement for handling custom models in `weights_only=True` mode.

### Alternatives

The current workaround involves setting `weights_only=False`, which allows repeatability but may pose a security risk due to the potential execution of arbitrary code during unpickling. However, this approach is not ideal for secure applications, especially when working with multiple models in neural architecture search.


### Additional context

Environment details:
- PyTorch version: 2.1.0
- Python version: 3.10
- OS: Ubuntu 22.04

Example code and detailed reproduction steps are provided below to illustrate the issue and expected vs. actual behavior.

```python
import torch
import numpy as np
import random
from my_custom_package import NeuralNetwork, Trainer  # Import custom model and trainer

# Set random seed for reproducibility
def set_random_seed(seed=42):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_random_seed(42)

# Initialize and train the model
model = NeuralNetwork(input_size=10, output_size=4, num_vars=40)
trainer = Trainer(model=model, lr=1e-4, batch_size=20, num_epochs=20)
trainer.fit(X_train, y_train, X_val, y_val)

# Save the trained model
trainer.save_model('neural_model.pth')

# Function to generate predictions
def generate_predictions(model, X_test):
    with torch.no_grad():
        X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
        return model(X_test_tensor)

# Generate predictions with the original model
predictions = generate_predictions(model, X_test)

# Load the model with weights_only=False and evaluate
loaded_model = torch.load('neural_model.pth', weights_only=False)
loaded_model.eval()
predictions_loaded_model = generate_predictions(loaded_model, X_test)

# Load the model again with default settings and evaluate
loaded_model_2 = torch.load('neural_model.pth')
loaded_model_2.eval()
predictions_loaded_model_2 = generate_predictions(loaded_model_2, X_test)

# Check consistency of predictions
print("All Same - trained and loaded:\t", np.allclose(predictions, predictions_loaded_model))
print("All Same - loaded once, loaded twice:\t", np.allclose(predictions_loaded_model_2, predictions_loaded_model_2))

```

Thank you for your time and assistance. Any guidance or recommendations would be greatly appreciated.


cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @mruberry @mikaylagawarecki