PackedQuantileMAELoss doesn't work properly #232

When i change loss_func in conf/finetune/model/moirai_1.1_T_large.yaml to uni2ts.loss.packed.PackedQuantileMAELoss, i get a error:

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
Seed set to 7
Number of windows in finetune: 91
Batch size for finetune: 128
Number of batches in a epoch: 0
/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/callbacks/model_checkpoint.py:751: Checkpoint directory /home/aptikal/laglilm/uni2ts/outputs/finetune/GIFTeval_M_moirai_1.1_R_large/jena_weather/D_short/checkpoints exists and is not empty.
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

| Name | Type | Params | Mode

0 | module | MoiraiModule | 310 M | train

310 M Trainable params
0 Non-trainable params
310 M Total params
1,243.882 Total estimated model params size (MB)
481 Modules in train mode
0 Modules in eval mode
Sanity Checking DataLoader 0: 0%| | 0/2 [00:00<?, ?it/s]Traceback (most recent call last):
File "/home/aptikal/laglilm/uni2ts/FinetuneGIFTEval/Finetune.py", line 180, in
run_training_for_datasets()
File "/home/aptikal/laglilm/uni2ts/FinetuneGIFTEval/Finetune.py", line 131, in run_training_for_datasets
main(cfg)
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/hydra/main.py", line 79, in decorated_main
return task_function(cfg_passthrough)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/uni2ts/cli/train.py", line 174, in main
trainer.fit(
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 560, in fit
call._call_and_handle_interrupt(
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 49, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 598, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1011, in _run
results = self._run_stage()
^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1053, in _run_stage
self._run_sanity_check()
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1082, in _run_sanity_check
val_loop.run()
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/loops/utilities.py", line 179, in _decorator
return loop_run(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 145, in run
self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 437, in _evaluation_step
output = call._call_strategy_hook(trainer, hook_name, *step_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 329, in _call_strategy_hook
output = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py", line 412, in validation_step
return self.lightning_module.validation_step(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/uni2ts/src/uni2ts/model/moirai/finetune.py", line 213, in validation_step
val_loss = self.hparams.loss_func(
^^^^^^^^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/uni2ts/src/uni2ts/loss/packed/_base.py", line 58, in call
loss = self._loss_func(
^^^^^^^^^^^^^^^^
File "/home/aptikal/laglilm/uni2ts/src/uni2ts/loss/packed/quantile.py", line 34, in _loss_func
quantile_levels = torch.tensor(self.quantile_levels, device=pred.device).view(
^^^^^^^^^^^
AttributeError: 'AffineTransformed' object has no attribute 'device'
/home/aptikal/laglilm/miniconda3/envs/aguni2ts/lib/python3.11/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
warnings.warn('resource_tracker: There appear to be %d '