import torch
import torch.nn as nn
from model import resnet18_cifar
import os

def create_demo_checkpoint():
    """
    Create a demo checkpoint for demonstration purposes.
    In a real deployment, you would use a properly trained model.
    """
    model = resnet18_cifar(num_classes=100, width=64)
    
    # Initialize with random weights (in practice, use trained weights)
    checkpoint = {
        'epoch': 100,
        'model_state': model.state_dict(),
        'best_acc1': 75.0,  # Example accuracy
        'args': {
            'width': 64,
            'num_classes': 100,
            'max_lr': 0.1,
            'weight_decay': 5e-4,
        }
    }
    
    # Save demo checkpoint
    os.makedirs('checkpoints', exist_ok=True)
    torch.save(checkpoint, 'checkpoints/demo_model.pth')
    print("Demo checkpoint created at checkpoints/demo_model.pth")
    
    return checkpoint

if __name__ == "__main__":
    create_demo_checkpoint()