class_docstr: Universal Transformer(https://arxiv.org/pdf/1807.03819.pdf) with ACT(https://arxiv.org/pdf/1603.08983.pdf).
code: def step_forward(self, nodes):
      x = nodes.data["x"]
      step = nodes.data["step"]
      pos = nodes.data["pos"]
      return {
          "x": self.pos_enc.dropout(
              x
              + self.pos_enc(pos.view(-1))
              + self.time_enc(step.view(-1))
          ),
          "step": step + 1,
      }
  
imports: from .attention import *
  from .layers import *
  from .functions import *
  from .embedding import *
  import dgl.function
  import torch
  import torch.nn.init
label_desc: Use image node_id 3 for calling the UTransformer obj's underlying member method code with example usage: obj.step_forward(nodes) and returns: dict