Fix bugs

Signed-off-by: Shanyu Zhao <shzhao@microsoft.com>