def get_model_attribute(model, attribute_name, cuda_device):
    """
        Getter function for both CPU and GPU.

        Parameters
        ____________________
        model: MultiTaskModel object,
        attribute_name: str

        Returns
        --------------------
        The attribute object from the model.
    """
    # maybe we should do (int, list)
    if isinstance(cuda_device, list):
        return getattr(model.module, attribute_name)
    else:
        return getattr(model, attribute_name)