I am wondering if there is a way to add a custom regression head(s) to unet model. I am thinking of replacing a classification_head with a regression model. Is it the correct way to tackle my problem?