text: import pytest
  import torch
  
  from src.models.model import create_model
  
  
  class TestClass:
      @pytest.mark.parametrize("output_shape, batch_size", [(10, 64), (5, 32), (59, 1)])
      def test_output_shape(self, output_shape, batch_size):
          model = create_model(784, output_shape, [256, 64])
          input = torch.rand(size=[batch_size, 28, 28])
          output = model(input)
          assert list(output.size()) == [batch_size, output_shape], \
              "The output of the model did not have the correct shape"
  
      def test_error_on_wrong_shape(self):
          model = create_model(784, 10, [256, 64])
          with pytest.raises(ValueError, match='Expected input to a 3D tensor'):
              model(torch.randn(1, 3))
          with pytest.raises(ValueError, match='Expected each sample to have shape 1, 28, 28'):
              model(torch.randn(1, 3, 5))
  