import torch


def _get_model_analysis_input(cfg, use_train_input):
    """
    Return a dummy input for model analysis with batch size 1. The input is
        used for analyzing the model (counting flops and activations etc.).
    Args:
        cfg (Config): the global config object.
        use_train_input (bool): if True, return the input for training. Otherwise,
            return the input for testing.

    Returns:
        Args: the input for model analysis.
    """
    rgb_dimension = 3
    if use_train_input:
        input_tensors = torch.rand(
            rgb_dimension,
            cfg.DATA.NUM_INPUT_FRAMES,
            cfg.DATA.TRAIN_CROP_SIZE,
            cfg.DATA.TRAIN_CROP_SIZE,
        )
    else:
        input_tensors = torch.rand(
            rgb_dimension,
            cfg.DATA.NUM_INPUT_FRAMES,
            cfg.DATA.TEST_CROP_SIZE,
            cfg.DATA.TEST_CROP_SIZE,
        )
    model_inputs = input_tensors.unsqueeze(0)
    if cfg.NUM_GPUS:
        model_inputs = model_inputs.cuda(non_blocking=True)
    inputs = {"video": model_inputs}
    return inputs