code: def fit(model_cls):
      train_iter = WikiText2(split="train")
      tokenizer = get_tokenizer("basic_english")
      vocab = build_vocab_from_iterator(
          map(tokenizer, train_iter), specials=["<unk>"]
      )
      vocab.set_default_index(vocab["<unk>"])
  
      def process_data(raw_text_iter):
          """Converts raw text into a flat Tensor."""
          data = [
              torch.tensor(vocab(tokenizer(item)), dtype=torch.long)
              for item in raw_text_iter
          ]
          return torch.cat(tuple(filter(lambda t: t.numel() > 0, data)))
  
      train_iter, val_iter, _ = WikiText2()
      train_data = process_data(train_iter)
      val_data = process_data(val_iter)
  
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
  
      def generate_batches(data, bsz):
          """Divides the data into bsz separate sequences."""
          seq_len = data.size(0) // bsz
          data = data[: seq_len * bsz]
          data = data.view(bsz, seq_len).t().contiguous()
          return data.to(device)
  
      batch_size = 20
      eval_batch_size = 10
      train_data = generate_batches(train_data, batch_size)
      val_data = generate_batches(val_data, eval_batch_size)
  
      seq_len = 35
  
      def get_seq(source, i):
          """
          Args:
              source: Tensor, with size [full_seq_len, batch_size]
              i: int
  
          Returns:
              tuple (data, target): data has size [seq_len, batch_size]
              and target has size [seq_len * batch_size]
          """
          part_len = min(seq_len, len(source) - 1 - i)
          data = source[i : i + part_len]
          target = source[i + 1 : i + 1 + part_len].reshape(-1)
          return data, target
  
      def generate_square_subsequent_mask(sz):
          """Generates an upper-triangular matrix of -inf, with zeros on diag."""
          return torch.triu(
              torch.ones(sz, sz) * float("-inf"), diagonal=1
          )
  
      model = model_cls().to(device)
      lr = 5.0
      optimizer = torch.optim.SGD(model.parameters(), lr=lr)
      scheduler = torch.optim.lr_scheduler.StepLR(
          optimizer, 1.0, gamma=0.95
      )
  
      def train(model):
          model.train()
          src_mask = generate_square_subsequent_mask(seq_len).to(device)
          for i in range(0, train_data.size(0) - 1, seq_len):
              data, target = get_seq(train_data, i)
              part_len = data.size(0)
              if part_len != seq_len:
                  src_mask = src_mask[:part_len, :part_len]
              output = model(data, src_mask)
              loss = F.cross_entropy(
                  output.view(-1, output.size(-1)), target
              )
  
              optimizer.zero_grad()
              loss.backward()
              torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
              optimizer.step()
  
      def evaluate(model, eval_data):
          model.eval()
          src_mask = generate_square_subsequent_mask(seq_len).to(device)
          total_loss = 0.0
          with torch.no_grad():
              for i in range(0, eval_data.size(0) - 1, seq_len):
                  data, target = get_seq(eval_data, i)
                  part_len = data.size(0)
                  if part_len != seq_len:
                      src_mask = src_mask[:part_len, :part_len]
                  output = model(data, src_mask)
                  output_flat = output.view(-1, output.size(-1))
                  total_loss += (
                      part_len
                      * F.cross_entropy(output_flat, target).item()
                  )
          return total_loss / (len(eval_data) - 1)
  
      best_val_loss = float("inf")
  
      for epoch in range(20):
          train(model)
          val_loss = evaluate(model, val_data)
          if val_loss < best_val_loss:
              best_val_loss = val_loss
          scheduler.step()
  
      best_val_ppl = math.exp(best_val_loss)
      nni.report_final_result(best_val_ppl)
  
imports: import torch
  import torch.nn.functional
  import nni.retiarii.nn.pytorch
  from nni.retiarii import model_wrapper
  import nni
  import nni.retiarii.strategy
  from nni.retiarii.evaluator import FunctionalEvaluator
  from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
  import math
  from torchtext.datasets import WikiText2
  from torchtext.data.utils import get_tokenizer
  from torchtext.vocab import build_vocab_from_iterator
label_desc: Use image node_id 1 for calling a global function with example usage: fit(model_cls) and returns: torch, data, data,  target, torch, unknown