Lightning-CLI use timm.create_model to initialize model in config file

implementation help
1

The Timm library offers a great variety of models, and with hydra/omegaconf one could write a config file like

model:
    _target_: timm.create_model
   model_name: resnet18

And be able to initiatlize a LightningModule, that takes model: nn.Module as an argument. I tried doing the same with LightningCLI but am not sure how to configure that correctly, because Import path timm.create_model does not correspond to a subclass of <class 'torch.nn.modules.module.Module'>. Here is the example:

from lightning import LightningModule, LightningDataModule
from lightning.pytorch.cli import LightningCLI, ArgsType
from lightning.pytorch.cli import OptimizerCallable

import os
import torch
from typing import Optional
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader


from lightning.pytorch.cli import OptimizerCallable, LightningCLI


class RandomDataset(Dataset):
    def __init__(self, size=5, num_samples=10):
        self.len = num_samples
        self.data = torch.randn(num_samples, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len
    
class RandomDataModule(LightningDataModule):
    def __init__(self, size=5, num_samples=10, batch_size=2):
        super().__init__()
        self.size = size
        self.num_samples = num_samples
        self.batch_size = batch_size

    def setup(self, stage=None):
        self.train_dataset = RandomDataset(self.size, self.num_samples)
        self.val_dataset = RandomDataset(self.size, self.num_samples)
        self.test_dataset = RandomDataset(self.size, self.num_samples)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size)

class BoringModel(LightningModule):
    def __init__(self, model: nn.Module, loss_fn: nn.Module = nn.CrossEntropyLoss(), optimizer: OptimizerCallable = torch.optim.Adam):
        super().__init__()

        self.model = model
        self.optimizer = optimizer
        self.loss_fn = loss_fn
    
    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def exclude_some_params(self):
        return self.model.network[0].parameters()

    def configure_optimizers(self):
        return self.optimizer(self.exclude_some_params())


def cli_main(args: ArgsType = None):
    LightningCLI(args=args, auto_configure_optimizers=False, parser_kwargs={"error_handler": None})

if __name__ == "__main__":
    cli_main(["fit", "--config", "example_yaml.yaml"])

And the config file:

trainer: # pytorch lightning trainer arguments
  max_epochs: 2

data: # datamodule arguments
  class_path: example.RandomDataModule
  init_args:
    size: 2

model:
  class_path: example.BoringModel
  init_args:
    model:
      class_path: timm.create_model
      init_args:
        model_name: "resnet18"
        num_classes: 2
        in_chans: 3
        pretrained: True
    optimizer:
      class_path: torch.optim.SGD
      init_args:
        lr: 0.003