Fixed import issues due to change in method signature