import torch


def interpolate_tensor(tensor, ratio: float = 1.0):
    """
    Upsample or Downsample tensor in time dimension
    :param tensor: (batch_size, n_timesteps, n_classes,...) or (batch_size, n_timesteps). Torch tensor.
    :param: ratio. If ratio > 1: upsample, ratio < 1: downsample # ratio = output rate/input rate
    :return: new_tensor (batch_size, n_timestepss*ratio, ...)
    """
    ratio = float(ratio)

    batch_size, n_input_frames = tensor.shape[0], tensor.shape[1]

    # compute indexes
    n_output_frames = int(round(n_input_frames * ratio))
    output_idx = torch.arange(n_output_frames)
    input_idx = torch.floor(output_idx / ratio).long()

    new_tensor = tensor[:, input_idx]

    return new_tensor