model = CropModel(num_classes=2)  # num_classes must be known beforehand

# Loading from checkpoint without needing to specify num_classes
model = CropModel.load_from_checkpoint("path/to/checkpoint.ckpt")

class CropModel(LightningModule):
    def __init__(self, num_classes=None, batch_size=4, num_workers=0, lr=0.0001, model=None, label_dict=None):
        super().__init__()
        self.num_classes = num_classes
        
        # Only initialize model if num_classes is provided
        if num_classes is not None:
            self._init_model(num_classes)
            
    def _init_model(self, num_classes):
        """Initialize model and metrics with given num_classes"""
        if getattr(self, 'model', None) is None:
            self.model = simple_resnet_50(num_classes=num_classes)
            
        self.accuracy = torchmetrics.Accuracy(average='none',
                                            num_classes=num_classes,
                                            task="multiclass")
        self.total_accuracy = torchmetrics.Accuracy(num_classes=num_classes,
                                                  task="multiclass")
        self.precision_metric = torchmetrics.Precision(num_classes=num_classes,
                                                     task="multiclass")
        self.metrics = torchmetrics.MetricCollection({
            "Class Accuracy": self.accuracy,
            "Accuracy": self.total_accuracy,
            "Precision": self.precision_metric
        })
            
    def on_load_checkpoint(self, checkpoint):
        if self.num_classes is None:
            # Extract num_classes from checkpoint
            self.num_classes = checkpoint['hyper_parameters']['num_classes']
            self._init_model(self.num_classes)