code: def forward(self, source_seq, target_seq):
      source_embeddings = self.embed_tokens(source_seq)
      target_embeddings = self.embed_tokens(target_seq)
  
      output = self.transformer(source_embeddings, target_embeddings)
  
      # Using integer key here, just like in Huggingface transformer lib
      return self.output_proj(output[0])
  
imports: from collections import OrderedDict
  from dataclasses import dataclass
  import tempfile
  import unittest
  import torch
  from torch import nn
  from fairscale.fair_dev.testing.testing import dist_init
  from fairscale.nn import FullyShardedDataParallel
  from fairscale.nn import auto_wrap, enable_wrap
label_desc: Use image node_id 3 for calling the TransformerWithLMHead obj's underlying member method code with example usage: obj.forward(source_seq, target_seq) and returns: self