Question Color mesh by height

First of all, Im a noob in shaders.

Ok, I have a procedurally generated terrain mesh, and I want it to have sand at the lower sides, and grass at higher ones.
I made so a script sets mesh.colors based on a gradient and in shader graph I made that it sets the albedo to the vertex colors. But I want to my terrain colors to not be smothed out, and I dont know how. I want to make something like in Sebastian Lagues terrain generation, where there is a line where a color ends and another starts. I found that with the vertex colors there are some places where the sand goes higher than another. The shader in Sebastians tutorial is a surface shader, and I think is not compatible with URP.

Any ideas?