def METHOD_NAME(module, input, output):
    class_name = str(module.__class__).split(".")[-1].split("'")[0]
    module_idx = len(summary)
    m_key = f"{class_name}-{module_idx + 1}"
    summary[m_key] = OrderedDict()
    # block_name refers to all layers that contains other layers
    if len(module._modules) != 0:
        summary[m_key]["block_name"] = class_name
    summary[m_key]["inference_time"] = np.round(timer.stop(), 3)
    timer.start()
    summary[m_key]["gpu_occupation"] = (round(torch.cuda.memory_allocated(0) / 1024**3, 2), "GB") if torch.cuda.is_available() else [0]
    summary[m_key]["gpu_cached_memory"] = (round(torch.cuda.memory_reserved(0) / 1024**3, 2), "GB") if torch.cuda.is_available() else [0]
    summary[m_key]["input_shape"], summary[m_key]["output_shape"] = get_input_output_shapes(batch_size=batch_size, input_dims=input, output_dims=output)
    params = 0
    if hasattr(module, "weight") and hasattr(module.weight, "size"):
        params += torch.prod(torch.LongTensor(list(module.weight.size())))
        summary[m_key]["trainable"] = module.weight.requires_grad
    if hasattr(module, "bias") and hasattr(module.bias, "size"):
        params += torch.prod(torch.LongTensor(list(module.bias.size())))
    summary[m_key]["nb_params"] = params