Guidance on how to do regression with ConvNets

Hello world, hope all is well,

Today I am trying to use my ConvNet skills and apply them to doing a regression problem. Unfortunately, there is little documentation compared to classification. Here (https://arxiv.org/pdf/1803.08450.pdf) I found the basic advice, “Besides classification,
ConvNets are also used to solve regression problems. In this
case, the softmax layer is commonly replaced with a fully
connected regression layer with linear or sigmoid activations.”

In my problem I have classes 0 through 4 and I want the algorithm to return a regression. After following the above advice all scores were between 1.8 and 4.1 - almost completely ignoring the 0 and 1 class. The results for classes 2,3,and 4 also look nonsensical. I am using MSELoss.

Any advice on how to start a problem like this correctly would be appreciated!
Thanks,

jjosh