Simplify the code

Simplify the code as @SpecLad suggested.