Fix build

Add back required code due to eab982b41c99e649c6e766edb0f857e34f3e0961