code: def _integration(dataloader, tmp_path, trainer_kwargs=None, **kwargs):
      train_dataloader = dataloader["train"]
      val_dataloader = dataloader["val"]
      test_dataloader = dataloader["test"]
  
      early_stop_callback = EarlyStopping(
          monitor="val_loss",
          min_delta=1e-4,
          patience=1,
          verbose=False,
          mode="min",
      )
  
      logger = TensorBoardLogger(tmp_path)
      if trainer_kwargs is None:
          trainer_kwargs = {}
      trainer = pl.Trainer(
          max_epochs=2,
          gradient_clip_val=0.1,
          callbacks=[early_stop_callback],
          enable_checkpointing=True,
          default_root_dir=tmp_path,
          limit_train_batches=2,
          limit_val_batches=2,
          limit_test_batches=2,
          logger=logger,
          **trainer_kwargs,
      )
  
      kwargs.setdefault("learning_rate", 0.15)
      kwargs.setdefault("weight_decay", 1e-2)
  
      net = NHiTS.from_dataset(
          train_dataloader.dataset,
          log_gradient_flow=True,
          log_interval=1000,
          hidden_size=8,
          **kwargs,
      )
      net.size()
      try:
          trainer.fit(
              net,
              train_dataloaders=train_dataloader,
              val_dataloaders=val_dataloader,
          )
          # todo: testing somehow disables grad computation even though it is explicitly turned on
          #       loss is calculated as "grad" for MQF2
          if not isinstance(net.loss, MQF2DistributionLoss):
              test_outputs = trainer.test(
                  net, dataloaders=test_dataloader
              )
              assert len(test_outputs) > 0
          # check loading
          net = NHiTS.load_from_checkpoint(
              trainer.checkpoint_callback.best_model_path
          )
  
          # check prediction
          net.predict(
              val_dataloader,
              fast_dev_run=True,
              return_index=True,
              return_decoder_lengths=True,
              trainer_kwargs=trainer_kwargs,
          )
      finally:
          shutil.rmtree(tmp_path, ignore_errors=True)
  
      net.predict(
          val_dataloader,
          fast_dev_run=True,
          return_index=True,
          return_decoder_lengths=True,
      )
  
file_path: pytorch-forecasting/tests/test_models/test_nhits.py
imports: import pickle
  import shutil
  import lightning.pytorch
  from lightning.pytorch.callbacks import EarlyStopping
  from lightning.pytorch.loggers import TensorBoardLogger
  import pytest
  from pytorch_forecasting.metrics import MQF2DistributionLoss, QuantileLoss
  from pytorch_forecasting.metrics.distributions import ImplicitQuantileNetworkDistributionLoss
  from pytorch_forecasting.models import NHiTS
label_desc: Use image node_id 1 for calling a global function with example usage: _integration(dataloader, tmp_path, trainer_kwargs) without return types