from pathlib import Path
import json


def load_train_state(path):
    """
    Load train_state parameters

    Parameters
    ----------
        path : str
            Directory path

    Returns
    -------
        train_state: Dict
            train_state parameters

    """

    file_name = Path(path, "train_state.json").as_posix()
    with open(file_name) as json_file:
        train_state = json.load(json_file)

    return train_state