[BUG] PLForecastingModule._calculate_metrics should call log_dict with batch_size

[youurayy]

Describe the bug
When using additional metrics with PyTorch Lightning forecasting models, due to PLForecastingModule._calculate_metrics (called from train_step() or validation_step()) not calling self.log_dict(...) with the batch_size= parameter, the Trainer (_ResultCollection.log() is complaining about ambiguity (when show_warnings=True is enabled on the model):

/opt/homebrew/Caskroom/miniconda/base/envs/tf/lib/python3.12/site-packages/pytorch_lightning/utilities/data.py:77: Trying to infer the batch_size from an ambiguous collection. The batch size we found is 32. To avoid any miscalculations, use self.log(..., batch_size=batch_size).

To Reproduce
Add a custom metric via torch_metrics=... model parameter (e.g. BlockRNNModel), and run .fit().

Expected behavior
The warnings should be gone.

System (please complete the following information):

• Python version: 3.12.3
• darts version: 0.29.0

Additional context
suggested code modification (tested):

def _calculate_metrics(self, output, target, metrics):
    if not len(metrics):
        return

    if self.likelihood:
        _metric = metrics(self.likelihood.sample(output), target)
    else:
        # If there's no likelihood, nr_params=1, and we need to squeeze out the
        # last dimension of model output, for properly computing the metric.
        _metric = metrics(output.squeeze(dim=-1), target)

    self.log_dict(
        _metric,
        on_epoch=True,
        on_step=False,
        logger=True,
        prog_bar=True,
        sync_dist=True,
        batch_size=target.shape[0],  # ADD THIS LINE
    )

[dennisbader]

Hi @youurayy and thanks for raising this issue.

Your suggested modification looks great! Would you like to contribute this?

[youurayy]

Hi @youurayy and thanks for raising this issue.

Your suggested modification looks great! Would you like to contribute this?

ah sorry @dennisbader I didn't check this in time. Thanks for fixing this!