Bug Fixes in new feature

Signed-off-by: Vishnu <vishnu.garg@shephertz.co.in>