code: def build_linear_or_identity(input_dim, output_dim):
      """
      Auxiliary method to return FC layer when input_dim != output_dim
      else return identity
      """
      if input_dim != output_dim:
          model = torch.nn.Linear(input_dim, output_dim)
      else:
          model = torch.nn.Identity()
  
      return model
  
imports: from typing import Dict, Optional
  import numpy
  import torch
  from pytorch_lightning import Trainer
  from nemo.collections.common.losses import NLLLoss
  from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTBottleneckModelConfig
  from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
  from nemo.core.classes.common import typecheck
  from nemo.utils import timers
label_desc: Use image node_id 1 for calling a global function with example usage: build_linear_or_identity(input_dim, output_dim) and returns: model