import torch
import pickle
import io

# First, let's create a simple PyTorch model
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)
    
    def forward(self, x):
        return self.linear(x)

# Create an instance of the model
model = SimpleModel()

# Now, let's optimize the model with torch.compile
# This should wrap it in an OptimizedModule
optimized_model = torch.compile(model)

# Check the type of the optimized model
print(f"Type of optimized model: {type(optimized_model)}")
print(f"Is it an OptimizedModule? {isinstance(optimized_model, torch._dynamo.OptimizedModule)}")

# Try to pickle the optimized model
try:
    # Use a BytesIO object to pickle in memory
    buffer = io.BytesIO()
    pickle.dump(optimized_model, buffer)
    buffer.seek(0)
    
    # Try to unpickle
    unpickled_model = pickle.load(buffer)
    
    print("Successfully pickled and unpickled the OptimizedModule!")
    print(f"Type of unpickled model: {type(unpickled_model)}")
    
    # Check if the unpickled model works
    test_input = torch.randn(2, 10)
    original_output = optimized_model(test_input)
    unpickled_output = unpickled_model(test_input)
    
    # Compare outputs
    print(f"Outputs match: {torch.allclose(original_output, unpickled_output)}")
    
except Exception as e:
    print(f"Failed to pickle/unpickle: {e}")
    print("torch._dynamo.OptimizedModule is NOT picklable")