code: def transpose_tensor(
      shape_tensor: torch.Tensor, input: torch.Tensor
  ) -> torch.Tensor:
      shape = len(shape_tensor.shape)
      assert shape in [2, 3], f"Invalid input shape {shape}"
      if shape == 2:
          return input
      else:
          return input.transpose(1, 2)
  
imports: import logging
  import math
  from typing import List, Optional
  import torch
  import torch.nn
  import torch.nn.init
  from reagent.core import types
  from reagent.models.base import ModelBase
  from reagent.models.residual_wrapper import ResidualWrapper
label_desc: Use image node_id 2 for calling a global function with example usage: transpose_tensor(shape_tensor, input) and returns: input, input