Encoding time of day as exogenous variable leads to error in GRIN model

[JanMaeg]

Hello,

I am trying to run imputation on the GRIN model. When I am trying to encode the timestamp as an additional exogenous variable with the following code, training the model fails. It does however work when using the SPIN imputation model.

covariates = {'u': dataset.datetime_encoded('day').values}

Stack trace of the failed training:

Traceback (most recent call last):
  File "Path\Lib\site-packages\tsl\experiment\experiment.py", line 169, in decorated_run_fn
    self.run_output = func(cfg)
                      ~~~~^^^^^
  File "Path\tryout_grin_exog.py", line 181, in run_imputation
    trainer.fit(imputer, datamodule=dm)
    ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\pytorch_lightning\trainer\trainer.py", line 560, in fit
    call._call_and_handle_interrupt(
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "Path\Lib\site-packages\pytorch_lightning\trainer\call.py", line 49, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "Path\Lib\site-packages\pytorch_lightning\trainer\trainer.py", line 598, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
    ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1011, in _run
    results = self._run_stage()
  File "Path\Lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1053, in _run_stage
    self._run_sanity_check()
    ~~~~~~~~~~~~~~~~~~~~~~^^
  File "Path\Lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1082, in _run_sanity_check
    val_loop.run()
    ~~~~~~~~~~~~^^
  File "Path\Lib\site-packages\pytorch_lightning\loops\utilities.py", line 179, in _decorator
    return loop_run(self, *args, **kwargs)
  File "Path\Lib\site-packages\pytorch_lightning\loops\evaluation_loop.py", line 145, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
    ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\pytorch_lightning\loops\evaluation_loop.py", line 437, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_args)
  File "Path\Lib\site-packages\pytorch_lightning\trainer\call.py", line 329, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "Path\Lib\site-packages\pytorch_lightning\strategies\strategy.py", line 412, in validation_step
    return self.lightning_module.validation_step(*args, **kwargs)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\tsl\engines\imputer.py", line 217, in validation_step
    y_hat, y, val_loss = self.shared_step(batch, batch.mask)
                         ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\tsl\engines\imputer.py", line 183, in shared_step
    y_hat = y_hat_loss = self.predict_batch(
                         ~~~~~~~~~~~~~~~~~~^
        batch, preprocess=False, postprocess=not self.scale_target)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\tsl\engines\predictor.py", line 285, in predict_batch
    y_hat = self.forward(**inputs, **forward_kwargs)
  File "Path\Lib\site-packages\tsl\engines\predictor.py", line 176, in forward
    return self.model(*args, **kwargs)
           ~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\torch\nn\modules\module.py", line 1879, in _call_impl
    return inner()
  File "Path\Lib\site-packages\torch\nn\modules\module.py", line 1827, in inner
    result = forward_call(*args, **kwargs)
  File "Path\Lib\site-packages\tsl\nn\models\stgn\grin_model.py", line 111, in forward
    fwd_out, fwd_pred, fwd_repr, _ = self.fwd_gril(x,
                                     ~~~~~~~~~~~~~^^^
                                                   edge_index,
                                                   ^^^^^^^^^^^
                                                   edge_weight,
                                                   ^^^^^^^^^^^^
                                                   mask=mask,
                                                   ^^^^^^^^^^
                                                   u=u)
                                                   ^^^^
  File "Path\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "Path\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "Path\Lib\site-packages\tsl\nn\layers\recurrent\grin.py", line 281, in forward
    inputs = torch.cat(inputs, dim=-1)  # x_hat_2 + mask + exogenous
RuntimeError: Tensors must have same number of dimensions: got 3 and 2

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

I re-created the error also running the provided imputation example from https://github.com/TorchSpatiotemporal/tsl/blob/main/examples/imputation/run_imputation_experiment.py

I am running the latest version of TSL torch-spatiotemporal==0.9.5.

Is there anything I need to configure different here when using GRIN?

Thanks in advance!