class_docstr: UNet model for segmentation task.
class_docstr_tok: ["UNet","model","for","segmentation","task","."]
code: def build_flax_model(self):
      return UNet(
          num_classes=self.dataset_meta_data["num_classes"],
          padding=self.config.model.get("padding", "SAME"),
          use_batch_norm=self.config.model.get("use_batch_norm", True),
          block_size=self.config.model.get(
              "block_size", (64, 128, 256, 512)
          ),
      )
  
imports: import functools
  from typing import Tuple
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.segmentation_model import SegmentationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.model_lib.layers import nn_ops
label_desc: Use image node_id 1 for calling the UNetSegmentationModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: UNet