Try something ?

After my PR, yours got conflict :P I think this can solve it. (Not sure but well ^^)