code: def create_model(
      config: ml_collections.ConfigDict, deterministic: bool
  ) -> nn.Module:
      """Creates a Flax model, as specified by the config."""
      if config.model == "GraphNet":
          return models.GraphNet(
              latent_size=config.latent_size,
              num_mlp_layers=config.num_mlp_layers,
              message_passing_steps=config.message_passing_steps,
              output_globals_size=config.num_classes,
              dropout_rate=config.dropout_rate,
              skip_connections=config.skip_connections,
              layer_norm=config.layer_norm,
              use_edge_model=config.use_edge_model,
              deterministic=deterministic,
          )
      if config.model == "GraphConvNet":
          return models.GraphConvNet(
              latent_size=config.latent_size,
              num_mlp_layers=config.num_mlp_layers,
              message_passing_steps=config.message_passing_steps,
              output_globals_size=config.num_classes,
              dropout_rate=config.dropout_rate,
              skip_connections=config.skip_connections,
              layer_norm=config.layer_norm,
              deterministic=deterministic,
          )
      raise ValueError(f"Unsupported model: {config.model}.")
  
imports: import os
  from typing import Any, Dict, Iterable, Tuple, Optional
  from absl import logging
  from clu import checkpoint
  from clu import metric_writers
  from clu import metrics
  from clu import parameter_overview
  from clu import periodic_actions
  import flax
  import flax.core
  import flax.linen
  from flax.training import train_state
  import jax
  import jax.numpy
  import jraph
  import ml_collections
  import numpy
  import optax
  import sklearn.metrics
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 1 for calling a global function with example usage: create_model(config, deterministic) and returns: models, models