class_docstr: Tests that checkpoints can be loaded.
class_docstr_tok: ["Tests","that","checkpoints","can","be","loaded","."]
docstr: Tests that real checkpoints can be loaded and used with the model.
docstr_tok: ["Tests","that","real","checkpoints","can","be","loaded","and","used","with","the","model","."]
code: def test_checkpoint_loading(self, config_module):
      """Tests that real checkpoints can be loaded and used with the model."""
      # We test the canonical checkpoint if there is one:
      try:
          config = config_module.get_config(
              init_mode="canonical_checkpoint"
          )
      except TypeError:
          config = config_module.get_config()
  
      module = models.TextZeroShotDetectionModule(
          body_configs=config.model.body,
          normalize=config.model.normalize,
          box_bias=config.model.box_bias,
      )
  
      # Parameter initialization:
      batch_size = 8
      img_size = config.dataset_configs.input_size
      num_queries = 10
      seq_len = config.dataset_configs.max_query_length
      images = jnp.ones((batch_size, img_size, img_size, 3))
      texts = jnp.ones(
          (batch_size, num_queries, seq_len), dtype=jnp.int32
      )
      init = functools.partial(module.init, train=False)
      init_params = jax.eval_shape(
          init, jax.random.PRNGKey(0), images, texts
      )
      params = module.bind({}).load(init_params, config.init_from)
  
      # Test running the model with the parameters:
      fn = functools.partial(module.apply, train=False)
      out = jax.eval_shape(fn, {"params": params}, images, texts)
  
      self.assertContainsSubset(
          {"pred_boxes", "pred_logits"}, out.keys()
      )
  
file_path: scenic/scenic/projects/owl_vit/tests/checkpoint_loading_test.py
imports: import functools
  import inspect
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  from scenic.projects.owl_vit import configs
  from scenic.projects.owl_vit import models
label_desc: Use image node_id 1 for calling the CheckpointLoadingTest obj's underlying member method code with example usage: obj.test_checkpoint_loading(config_module) without return types