Fix bugs 

Signed-off-by: Pengfei Zhang <pzhang11@stevens.edu>