docstr: Tests creating model.
docstr_tok: ["Tests","creating","model","."]
code: def test_create_model(self):
      """Tests creating model."""
      model = train.create_model(
          model_cls=models._ResNet1, half_precision=False
      )  # pylint: disable=protected-access
      params, batch_stats = train.initialized(random.key(0), 224, model)
      variables = {"params": params, "batch_stats": batch_stats}
      x = random.normal(random.key(1), (8, 224, 224, 3))
      y = model.apply(variables, x, train=False)
      self.assertEqual(y.shape, (8, 1000))
  
file_path: flax/examples/imagenet/train_test.py
imports: import pathlib
  import tempfile
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import random
  import tensorflow
  import tensorflow_datasets
  import models
  import train
  from configs import default
label_desc: Use image node_id 2 for calling the TrainTest obj's underlying member method code with example usage: obj.test_create_model() without return types