Fine-tuning Denoising Diffusion Probabilistic Model

Hi, I have pretrained Denoising Diffusion Probabilistic Model (DDPM) for image generation and I want to fine-tune it for image segmentation. I am using the official implementation from GitHub - lucidrains/denoising-diffusion-pytorch: Implementation of Denoising Diffusion Probabilistic Model in Pytorch.

Does anyone have experience with extracting weights from the U-Net model? The last layer of this implemented U-Unet has only one value as the output but I want output as the tensor with shape (batch, dim, img_shape, img_shape).

Thank you for the answers.