code: def test_train(data_format, grayscale, tmp_path):
      # TODO: Test for grayscale = False and different size.
      # data = create_dataset(
      #     data_format=data_format, size=28, grayscale=True
      # )
      # wrapper = TorchImageClassificationWrapper("simple-cnn")
      data = create_dataset(
          data_format=data_format,
          grayscale=grayscale,
          tmp_path=tmp_path,
      )
      wrapper = TorchImageClassificationWrapper(
          "resnet18", {}, tmp_path
      )
  
      # TODO: Test both resnet18 and simple-cnn with a few different configurations of
      #   data.
      # TODO: Test with and without val/test data.
      # TODO: Check that something was written to writer and experiment.
      wrapper._train(
          train_data=data,
          val_data=data,
          test_data=data,
          writer=SummaryWriter(write_to_disk=False),
          experiment=DummyExperiment(),
          dry_run=True,  # True,
      )
  
      assert isinstance(wrapper.model, nn.Module)
      assert (tmp_path / "model.pt").exists()
  
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 4 for calling a global function with example usage: test_train(data_format, grayscale, tmp_path) without return types