from typing import List
from typing import Dict
import torch
from typing import Iterable


def _concat_step_output(outputs: List[Dict[str, torch.Tensor]], dim: int = 0, skip: Iterable = ('loss',))\
        -> Dict[str, torch.Tensor]:
    """Concatenate step output into single batch tensor.

    Args:
        outputs: Outputs from train/validation/test_step.
        dim: Dimension to be concatenated.
        skip: Keys to be skipped in outputs.
    """
    data = {}
    for output in outputs:
        for key, var in output.items():
            if key in skip:
                continue

            if key in data:
                data[key].append(var)
            else:
                data[key] = [var]

    for key in data:
        data[key] = torch.cat(data[key], dim=dim).clone().detach().cpu()

    return data