Added check in both functions