code: def test_NTA_V1(n: int, last_n_trigger: bool):
      model = initialize_model()
      params = model.collect_params()
      loss_list = [5, 4, 3, 2, 3, 3, 3, 3, 3, 3, 3]
      avg_strategy = NTA(
          epochs=len(loss_list),
          n=n,
          last_n_trigger=last_n_trigger,
          fallback_alpha=0.0,
      )
      for i, loss in enumerate(loss_list):
          for k, v in params.items():
              for arr in v.list_data():
                  arr[:] = i
          avg_strategy.update_average_trigger(metric=loss)
          avg_strategy.apply(model)
      # nothing is cached yet, thus load_cached_model won't change anything
      # test cached model
      avg_strategy.load_cached_model(model)
      for k, v in params.items():
          for arr in v.list_data():
              # the last model should have 10 in all coordinates
              assert mx.nd.norm(arr - 10).asscalar() < 1e-30
      # test averaged model
      avg_strategy.load_averaged_model(model)
      len_limit = (len(loss_list) - 1) if last_n_trigger else 6
      if n <= 0 or n > len_limit:
          # average never happens, model is not changed
          for k, v in params.items():
              for arr in v.list_data():
                  # the last model should have 10 in all coordinates
                  assert mx.nd.norm(arr - 10).asscalar() < 1e-30
      else:
          for k, v in params.items():
              for arr in v.list_data():
                  if last_n_trigger:
                      # takes the average once the loss increases, no matter what n is taken
                      # (the first n iterations are ignored)
                      val = 7 if n <= 4 else ((n + 10) / 2.0)
                  else:
                      # takes the average on the last 7-n iterations
                      val = (4 + n + 10) / 2.0
                  assert mx.nd.norm(arr - val).asscalar() < 1e-30
      # test cached model
      avg_strategy.load_cached_model(model)
      for k, v in params.items():
          for arr in v.list_data():
              # the last model should have 10 in all coordinates
              assert mx.nd.norm(arr - 10).asscalar() < 1e-30
  
file_path: gluonts/test/mx/trainer/test_model_iteration_averaging.py
imports: import math
  import mxnet
  import mxnet.gluon.nn
  import numpy
  import pandas
  import pytest
  from gluonts.dataset.common import ListDataset
  from gluonts.mx import SimpleFeedForwardEstimator
  from gluonts.mx.trainer import Trainer
  from gluonts.mx.trainer.model_iteration_averaging import NTA, Alpha_Suffix
label_desc: Use image node_id 2 for calling a global function with example usage: test_NTA_V1(n, last_n_trigger) without return types