import torch


def load_model(model_path, map_location=None):
    """
    Wrapper function for `for safely loading models where certain new attributes of the model class are not present.
    E.g. "requires_stress" for computing the stress tensor.

    Args:
        model_path (str): Path to the model file.
        map_location (torch.device): Device where the model should be loaded to.

    Returns:
        :class:`schnetpack.atomistic.AtomisticModel`: Loaded SchNetPack model.

    """
    model = torch.load(model_path, map_location=map_location)

    # Check for data parallel models
    if hasattr(model, "module"):
        model_module = model.module
        output_modules = model.module.output_modules
    else:
        model_module = model
        output_modules = model.output_modules

    # Set stress tensor attribute if not present
    if not hasattr(model_module, "requires_stress"):
        model_module.requires_stress = False
        for module in output_modules:
            module.stress = None

    return model