This needs to be abstracted in a separate function, but in the meantime
we fix the issue in both places.

Signed-off-by: Iustin Pop <iustin@google.com>