from typing import Dict
import torch
from typing import List


def collect_flow_outputs_by_suffix(
    output_dict: Dict[str, torch.Tensor], suffix: str
) -> List[torch.Tensor]:
    """Return output_dict outputs specified by suffix, ordered by sorted flow_name."""
    return [
        output_dict[flow_name]
        for flow_name in sorted(output_dict.keys())
        if flow_name.endswith(suffix)
    ]