Fixed not calling correct function