code: def __init__(
      self,
      bottom_model: t.nn.Module,
      agglayer_arg: Union[
          StdAggLayerArgument, FedPassArgument, SSHEArgument
      ] = None,
      ctx: Context = None,
  ):
      super().__init__()
  
      assert isinstance(
          bottom_model, t.nn.Module
      ), "bottom model should be a torch nn.Module"
      self._bottom_model = bottom_model
      # cached variables
      self._bottom_fw = None  # for backward usage
      # ctx
      self._ctx = None
      self._agg_layer = None
      self._fake_loss = None
      self.setup(ctx=ctx, agglayer_arg=agglayer_arg)
  
imports: import torch
  import torch
  from fate.arch import Context
  from fate.ml.nn.model_zoo.agg_layer.agg_layer import AggLayerGuest, AggLayerHost
  from fate.ml.nn.model_zoo.agg_layer.fedpass.agg_layer import FedPassAggLayerGuest, FedPassAggLayerHost, get_model
  from fate.ml.nn.model_zoo.agg_layer.sshe.agg_layer import SSHEAggLayerHost, SSHEAggLayerGuest
  from typing import Any, Dict, List, Union, Callable, Literal
  from dataclasses import dataclass, fields
  from enum import Enum
  import logging
  from torch import device
label_desc: Use image node_id 1 to create a new HeteroNNModelHost object from inherited base classes: HeteroNNModelBase with example: obj = HeteroNNModelHost(bottom_model, agglayer_arg, ctx)