Training NN to map quantities into a grid

Question

I created a simple NN (multi-layer percetpron acutally) in pytorch that takes a vector of quantities as an input and outputs an assignment map so that the products with the biggest quantities are placed the closest to the origin of the 2d map (the manhattan distance is considered).

For an input of shape [5] (5 quantities), the output will be a tensor of shape [5+1,grid_height,grid_width] with probabilities of each product to be in each location, and the probabilities of no product to be at each location (hence the +1). Then, with an argmax, I am able to retrieve the final assignment.

For example, the input [5,2,3] should give me the result (after argmax) : [[1,1,1,2] [1,2,3,0] [2,3,0,0] [0,0,0,0]]

For now this is a very simple goal that could be done with a little algorithm but I want later to increase the amount of constraints.

The problem is that with very little quantities, it works almost perfectly, but when they are bigger, the input quantities are not respected, even if the highest quantities are well-placed near the origin.

I tried to add a constraint to my loss so that it penalizes the outputs that don't match the input's quantities, but it doesn't seem to solve the problem entirely.

I then tried to preprocess the data by sorting the input vector so that my NN can learn more efficiently, and it's my current solution but it is not very scalable to higher quantities or bigger input vectors.

What could be done to improve my system?