def predict_class(data, model, model_name: str):
    """
    Returns predicted classes for the given dataset and model.

    :param data: Training or test data.
    :param model: Trained model.
    :param model_name: Name of the trained model.
    :return: Predicted classes
    """
    if model_name == "MLP":
        output_activation = model.get_config()["layers"][-1]["config"]["activation"]
        if output_activation == "softmax":
            probs = model.predict(data)
            preds = probs.argmax(axis=-1)
        else:
            preds = model.predict_classes(data)
    else:
        preds = model.predict(data).astype("float64")

    return preds