Fix as per review comments

Change-Id: Id23bf8f8610f01dadcd5c137bf2a3b1888f4adb2