* changes * fix * fix for vars too * changes * fix tests Co-authored-by: Abubakar Abid <abubakar@huggingface.co>