code: def test_prediction_with_dataloder_raw(
      data_with_covariates, tmp_path
  ):
      # tests correct concatenation of raw output
      test_data = data_with_covariates.copy()
      np.random.seed(2)
      test_data = test_data.sample(frac=0.5)
  
      dataset = TimeSeriesDataSet(
          test_data,
          time_idx="time_idx",
          max_encoder_length=8,
          max_prediction_length=10,
          min_prediction_length=1,
          min_encoder_length=1,
          target="volume",
          group_ids=["agency", "sku"],
          constant_fill_strategy=dict(volume=0.0),
          allow_missing_timesteps=True,
          time_varying_unknown_reals=["volume"],
          time_varying_known_reals=["time_idx"],
          target_normalizer=GroupNormalizer(groups=["agency", "sku"]),
      )
  
      net = TemporalFusionTransformer.from_dataset(
          dataset,
          learning_rate=1e-6,
          hidden_size=4,
          attention_head_size=1,
          dropout=0.2,
          hidden_continuous_size=2,
          log_interval=1,
          log_val_interval=1,
          log_gradient_flow=True,
      )
      logger = TensorBoardLogger(tmp_path)
      trainer = pl.Trainer(
          max_epochs=1, gradient_clip_val=1e-6, logger=logger
      )
      trainer.fit(
          net,
          train_dataloaders=dataset.to_dataloader(
              batch_size=4, num_workers=0
          ),
      )
  
      # choose small batch size to provoke issue
      res = net.predict(
          dataset.to_dataloader(batch_size=2, num_workers=0), mode="raw"
      )
      # check that interpretation works
      net.interpret_output(res)["attention"]
      assert net.interpret_output(res.iget(slice(1)))[
          "attention"
      ].size() == torch.Size((1, net.hparams.max_encoder_length))
  
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 14 for calling a global function with example usage: test_prediction_with_dataloder_raw(data_with_covariates, tmp_path) without return types