def typecheck_hook(model, *, in_dtype=None, out_dtype=None):
    """
    Check that the model must be called with the given input/output dtype
    """
    if not isinstance(in_dtype, set):
        in_dtype = {in_dtype}
    if not isinstance(out_dtype, set):
        out_dtype = {out_dtype}

    def flatten(x):
        if isinstance(x, torch.Tensor):
            return [x]
        if isinstance(x, (list, tuple)):
            return list(itertools.chain(*[flatten(t) for t in x]))
        if isinstance(x, dict):
            return flatten(list(x.values()))
        return []

    def hook(module, input, output):
        if in_dtype is not None:
            dtypes = {x.dtype for x in flatten(input)}
            assert (
                dtypes == in_dtype
            ), f"Expected input dtype of {type(module)} is {in_dtype}. Got {dtypes} instead!"

        if out_dtype is not None:
            dtypes = {x.dtype for x in flatten(output)}
            assert (
                dtypes == out_dtype
            ), f"Expected output dtype of {type(module)} is {out_dtype}. Got {dtypes} instead!"

    with model.register_forward_hook(hook):
        yield

    Check that the model must be called with the given input/output dtype
    
video_chat_text/video_chat_with_ChatGPT/models/grit_src/third_party/CenterNet2/tests/modeling/test_model_e2e.py