Repository: Lightning-AI/pytorch-lightning
File: src/lightning/pytorch/core/module.py (Python)

Code before:
                self.log_dict({'val_loss': loss, 'val_acc': val_acc})

        If you pass in multiple val dataloaders, :meth:`validation_step` will have an additional argument. We recommend
        setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

        .. code-block:: python

            # CASE 2: multiple validation dataloaders
            def validation_step(self, batch, batch_idx, dataloader_idx=0):
                # dataloader_idx tells you which dataset this is.
                x, y = batch

                # implement your own
                out = self(x)

                if dataloader_idx == 0:
                    loss = self.loss0(out, y)
                else:
                    loss = self.loss1(out, y)

                # calculate acc
                labels_hat = torch.argmax(out, dim=1)
                acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

                # log the outputs separately for each dataloader
                self.log_dict({f"val_loss_{dataloader_idx}": loss, "val_acc_{dataloader_idx}": acc})

        Note:
            If you don't need to validate you don't need to implement this method.

        Note:
            When the :meth:`validation_step` is called, the model has been put in eval mode
            and PyTorch gradients have been disabled. At the end of validation,
            the model goes back to training mode and gradients are enabled.

        """

    def test_step(self, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
        r"""Operates on a single batch of data from the test set. In this step you'd normally generate examples or
        calculate anything of interest such as accuracy.

        Args:
            batch: The output of your data iterable, normally a :class:`~torch.utils.data.DataLoader`.
            batch_idx: The index of this batch.
            dataloader_idx: The index of the dataloader that produced this batch.
                (only if multiple dataloaders used)

        Return:
            - :class:`~torch.Tensor` - The loss tensor
            - ``dict`` - A dictionary. Can include any keys, but must include the key ``'loss'``.
            - ``None`` - Skip to the next batch.

Review comment (suggestion):
```suggestion
                self.log_dict({f"val_loss_{dataloader_idx}": loss, f"val_acc_{dataloader_idx}": acc})
```

Code after:
                self.log_dict({'val_loss': loss, 'val_acc': val_acc})

        If you pass in multiple val dataloaders, :meth:`validation_step` will have an additional argument. We recommend
        setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

        .. code-block:: python

            # CASE 2: multiple validation dataloaders
            def validation_step(self, batch, batch_idx, dataloader_idx=0):
                # dataloader_idx tells you which dataset this is.
                x, y = batch

                # implement your own
                out = self(x)

                if dataloader_idx == 0:
                    loss = self.loss0(out, y)
                else:
                    loss = self.loss1(out, y)

                # calculate acc
                labels_hat = torch.argmax(out, dim=1)
                acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

                # log the outputs separately for each dataloader
                self.log_dict({f"val_loss_{dataloader_idx}": loss, f"val_acc_{dataloader_idx}": acc})

        Note:
            If you don't need to validate you don't need to implement this method.

        Note:
            When the :meth:`validation_step` is called, the model has been put in eval mode
            and PyTorch gradients have been disabled. At the end of validation,
            the model goes back to training mode and gradients are enabled.

        """

    def test_step(self, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
        r"""Operates on a single batch of data from the test set. In this step you'd normally generate examples or
        calculate anything of interest such as accuracy.

        Args:
            batch: The output of your data iterable, normally a :class:`~torch.utils.data.DataLoader`.
            batch_idx: The index of this batch.
            dataloader_idx: The index of the dataloader that produced this batch.
                (only if multiple dataloaders used)

        Return:
            - :class:`~torch.Tensor` - The loss tensor
            - ``dict`` - A dictionary. Can include any keys, but must include the key ``'loss'``.
            - ``None`` - Skip to the next batch.