def module_tracker(fwd_hook_func):
 
    """
    Wrapper for tracking the layers throughout the forward pass.

    Arguments
    ---------
        
        fwd_hook_func : function
            Forward hook function to be wrapped.

    Returns
    -------

        function :
            Wrapped hook function

    """

    def hook_wrapper(layer, *args):
        return fwd_hook_func(layer, *args)

    return hook_wrapper