class_docstr: Simple CNN model for segmentation task.
class_docstr_tok: ["Simple","CNN","model","for","segmentation","task","."]
code: def build_flax_model(self) -> nn.Module:
      model_dtype = getattr(
          jnp, self.config.get("model_dtype_str", "float32")
      )
      return SimpleCNN(
          num_outputs=self.dataset_meta_data["num_classes"],
          num_filters=self.config.num_filters,
          kernel_sizes=self.config.kernel_sizes,
          use_bias=self.config.get("use_bias", None),
          output_projection_type="keep_dims",
          dtype=model_dtype,
      )
  
file_path: scenic/scenic/projects/baselines/simple_cnn.py
imports: from typing import Iterable, Callable, Sequence, Optional, Union
  import flax.linen
  from jax.nn import initializers
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.base_models.segmentation_model import SegmentationModel
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the SimpleCNNSegmentationModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: SimpleCNN