code: def _integration(
      dataloader, tmp_path, loss=None, 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",
      )
  
      # check training
      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,
      )
      # test monotone constraints automatically
      if "discount_in_percent" in train_dataloader.dataset.reals:
          monotone_constaints = {"discount_in_percent": +1}
          cuda_context = torch.backends.cudnn.flags(enabled=False)
      else:
          monotone_constaints = {}
          cuda_context = nullcontext()
  
      kwargs.setdefault("learning_rate", 0.15)
  
      with cuda_context:
          if loss is not None:
              pass
          elif isinstance(
              train_dataloader.dataset.target_normalizer,
              NaNLabelEncoder,
          ):
              loss = CrossEntropy()
          elif isinstance(
              train_dataloader.dataset.target_normalizer,
              MultiNormalizer,
          ):
              loss = MultiLoss(
                  [
                      CrossEntropy()
                      if isinstance(normalizer, NaNLabelEncoder)
                      else QuantileLoss()
                      for normalizer in train_dataloader.dataset.target_normalizer.normalizers
                  ]
              )
          else:
              loss = QuantileLoss()
          net = TemporalFusionTransformer.from_dataset(
              train_dataloader.dataset,
              hidden_size=2,
              hidden_continuous_size=2,
              attention_head_size=1,
              dropout=0.2,
              loss=loss,
              log_interval=5,
              log_val_interval=1,
              log_gradient_flow=True,
              monotone_constaints=monotone_constaints,
              **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 = TemporalFusionTransformer.load_from_checkpoint(
                  trainer.checkpoint_callback.best_model_path
              )
  
              # check prediction
              predictions = net.predict(
                  val_dataloader,
                  return_index=True,
                  return_x=True,
                  return_y=True,
                  fast_dev_run=True,
                  trainer_kwargs=trainer_kwargs,
              )
              pred_len = len(predictions.index)
  
              # check that output is of correct shape
              def check(x):
                  if isinstance(x, (tuple, list)):
                      for xi in x:
                          check(xi)
                  elif isinstance(x, dict):
                      for xi in x.values():
                          check(xi)
                  else:
                      assert (
                          pred_len == x.shape[0]
                      ), "first dimension should be prediction length"
  
              check(predictions.output)
              if isinstance(predictions.output, torch.Tensor):
                  assert (
                      predictions.output.ndim == 2
                  ), "shape of predictions should be batch_size x timesteps"
              else:
                  assert all(
                      p.ndim == 2 for p in predictions.output
                  ), "shape of predictions should be batch_size x timesteps"
              check(predictions.x)
              check(predictions.index)
  
              # predict raw
              net.predict(
                  val_dataloader,
                  return_index=True,
                  return_x=True,
                  fast_dev_run=True,
                  mode="raw",
                  trainer_kwargs=trainer_kwargs,
              )
  
          finally:
              shutil.rmtree(tmp_path, ignore_errors=True)
  
imports: import pickle
  import shutil
  import sys
  import lightning.pytorch
  from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
  from lightning.pytorch.loggers import TensorBoardLogger
  import numpy
  import pytest
  import torch
  from pytorch_forecasting import TimeSeriesDataSet
  from pytorch_forecasting.data import NaNLabelEncoder
  from pytorch_forecasting.data.encoders import GroupNormalizer, MultiNormalizer
  from pytorch_forecasting.metrics import CrossEntropy, MQF2DistributionLoss, MultiLoss, NegativeBinomialDistributionLoss, PoissonLoss, QuantileLoss, TweedieLoss
  from pytorch_forecasting.models import TemporalFusionTransformer
  from pytorch_forecasting.models.temporal_fusion_transformer.tuning import optimize_hyperparameters
  from test_models.conftest import make_dataloaders
label_desc: Use image node_id 5 for calling a global function with example usage: _integration(dataloader, tmp_path, loss, trainer_kwargs) without return types