Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/core/step_result.py (Python)

Code before:
    ):
        # set the meta for the item
        meta_value = value
        meta = dict(
            prog_bar=prog_bar,
            logger=logger,
            on_step=on_step,
            on_epoch=on_epoch,
            reduce_fx=reduce_fx,
            value=meta_value,
            tbptt_reduce_fx=tbptt_reduce_fx,
            tbptt_pad_token=tbptt_pad_token,
            forked=forked
        )

        self['meta'][name] = meta

        # track whether any input requires reduction on epoch end
        _internal = self['meta']['_internal']
        _internal['_reduce_on_epoch'] = max(_internal['_reduce_on_epoch'], on_epoch)

    def track_batch_size(self, batch):
        batch_size = self.unpack_batch_size(batch)
        meta = self['meta']
        meta['_internal']['batch_sizes'].append(batch_size)

    def get_batch_sizes(self):
        meta = self['meta']
        return torch.tensor(meta['_internal']['batch_sizes'])

    def get_callback_metrics(self) -> dict:
        result = {'early_stop_on': self.early_stop_on, 'checkpoint_on': self.checkpoint_on}

        return result

    def get_batch_log_metrics(self, include_forked_originals=True) -> dict:
        """
        Gets the metrics to log at the end of the batch step

        """
        result = {}

        meta = self['meta']
        for k, options in meta.items():
            if k == '_internal':
                continue

            if options['forked'] and not include_forked_originals:
                continue

            if options['logger'] and options['on_step']:

Review comment (performance):
@williamFalcon I'm not too sure about the correctness of setting this to 1, or of using `len(sample)` in `unpack_batch_size()` when `sample` is a string. This way, the batch size will always be 1, even we have a list of strings with `batch_size>1` coming from the dataloader. You can try to add a `batch_size` to the dataloader in the test, and `unpack_batch_size` will still indicate 1 as the batch size with the current patch.

Code after:
    ):
        # set the meta for the item
        meta_value = value
        meta = dict(
            prog_bar=prog_bar,
            logger=logger,
            on_step=on_step,
            on_epoch=on_epoch,
            reduce_fx=reduce_fx,
            value=meta_value,
            tbptt_reduce_fx=tbptt_reduce_fx,
            tbptt_pad_token=tbptt_pad_token,
            forked=forked
        )

        self['meta'][name] = meta

        # track whether any input requires reduction on epoch end
        _internal = self['meta']['_internal']
        _internal['_reduce_on_epoch'] = max(_internal['_reduce_on_epoch'], on_epoch)

    def track_batch_size(self, batch):
        try:
            batch_size = self.unpack_batch_size(batch)
        except RecursionError as re:
            batch_size = 1

        meta = self['meta']
        meta['_internal']['batch_sizes'].append(batch_size)

    def get_batch_sizes(self):
        meta = self['meta']
        return torch.tensor(meta['_internal']['batch_sizes'])

    def get_callback_metrics(self) -> dict:
        result = {'early_stop_on': self.early_stop_on, 'checkpoint_on': self.checkpoint_on}

        return result

    def get_batch_log_metrics(self, include_forked_originals=True) -> dict:
        """
        Gets the metrics to log at the end of the batch step

        """
        result = {}

        meta = self['meta']
        for k, options in meta.items():
            if k == '_internal':
                continue