Prediction integer-valued function

I would like to create a model with pytorch_geometric, which outputs on each node, a function from the segment [0,1] to the integers. This function is supposed to count a number of neighboring edges at a given time, so it has integer values, and it is locally constant.

I wonder how to efficiently implement such an output, and which loss function to use. If I merely output a vector that gives the outputs over a sample of [0,1], I am not using the fact that my function will be locally constant, so I am working in too much generality.
Any ideas?