import torch


def load_AE(ModelClass, path, **kwargs):
    """Loads an encoder and decoder"""
    model = ModelClass(**kwargs)
    model.load_state_dict(torch.load(path))
    #model.eval()
    encoder = model.encode
    decoder = model.decode

    return encoder, decoder