Bug fix and minor improvements

Change-Id: I83ca36fc94a55c155aa6a863f619f1d38396c134