import torch


def max_along_time(outputs, lengths):
    """ Get maximum responses from RNN outputs along time axis
    :param outputs: (B, T, D)
    :param lengths: (B, )
    :return: (B, D)
    """
    outputs = [outputs[i, :int(lengths[i]), :].max(dim=0)[0]
               for i in range(len(lengths))]

    return torch.stack(outputs, dim=0)