Bug fix

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