code: def __init__(self, config: StageConfig) -> None:
      super().__init__()
      ch_in = config.in_channels
      self.downsample = (
          ConvNormAct(ch_in, ch_in, 3, 2, "none", ch_in)
          if config.downsample
          else None
      )
      self.blocks = nn.Sequential(
          HGBlock(ch_in, config, False),
          *[
              HGBlock(config.out_channels, config, True)
              for _ in range(config.num_blocks - 1)
          ],
      )
  
imports: from __future__ import annotations
  from typing import NamedTuple
  from torch import nn
  from kornia.contrib.models.common import ConvNormAct
  from kornia.core import Module, Tensor, concatenate
  from kornia.core.check import KORNIA_CHECK
label_desc: Use image node_id 1 to create a new HGStage object from inherited base classes: nn with example: obj = HGStage(config)