code: def test_create_model(tmp_path):
      # resnet18
      wrapper = TorchImageClassificationWrapper(
          "resnet18", {"num_classes": 10}, tmp_path
      )
      wrapper._create_model(num_classes=10)
      assert isinstance(wrapper.model, nn.Module)
      assert wrapper.model.fc.out_features == 10
      # pretrained=True is not tested here because it takes too long
  
      # simple-cnn
      wrapper = TorchImageClassificationWrapper(
          "simple-cnn", {}, tmp_path
      )
      wrapper._create_model(num_classes=10)
      assert isinstance(wrapper.model, SimpleCnn)
  
file_path: traintool/tests/image_classification/test_torch_models.py
imports: import comet_ml
  import pytest
  import numpy
  import torch.nn
  from torch import optim
  from tensorboardX import SummaryWriter
  from conftest import create_dataset, create_image
  from traintool.image_classification.torch_models import TorchImageClassificationWrapper, SimpleCnn
  from traintool.utils import DummyExperiment
label_desc: Use image node_id 2 for calling a global function with example usage: test_create_model(tmp_path) without return types