# Using ViTMAEModel as an encoder for a UNet decoder for semantic segmentation

**URL:** <https://discuss.huggingface.co/t/using-vitmaemodel-as-an-encoder-for-a-unet-decoder-for-semantic-segmentation/91222>\
**Category:** 🤗Transformers\
**Created:** [June 9, 2024, 11:22pm UTC](https://discuss.huggingface.co/t/using-vitmaemodel-as-an-encoder-for-a-unet-decoder-for-semantic-segmentation/91222 "2024-06-09T23:22:37Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Mayolo](https://avatars.discourse-cdn.com/v4/letter/m/71c47a/32.png) [@Mayolo](https://discuss.huggingface.co/u/Mayolo)\
**Post date:** [June 9, 2024, 11:22pm UTC](https://discuss.huggingface.co/t/using-vitmaemodel-as-an-encoder-for-a-unet-decoder-for-semantic-segmentation/91222/1 "2024-06-09T23:22:37Z")

</div>

I want to train a Masked Autoencoder (ViTMAE) on a downstream task. The task is semantic segmentation and I have been trying to use the ViTMAEModel function as the encoder. I just don’t know what to do to be able to properly pass the last hidden state and all the other states into the decoder. Any ideas?

class UNetDecoder(nn.Module):  
def **init** (self,config):  
super(UNetDecoder, self). **init** ()  
#self.config = config

```
    self.up1 = nn.ConvTranspose2d(in_channels=768,out_channels=512,kernel_size=2,stride=2)
    self.up2 = nn.ConvTranspose2d(512,256,kernel_size=2,stride=2)
    self.up3 = nn.ConvTranspose2d(256,128,kernel_size=2,stride=2)
    self.up4 = nn.ConvTranspose2d(128,64,kernel_size=2,stride=2)
    self.out_conv = nn.Conv2d(64, config.num_channels, kernel_size=2)

def forward(self, skip_connection):
    
    x = self.up1(skip_connection[-1])
    x = torch.cat([x,skip_connection[-2]])
    
    x = self.up2(x)
    x = torch.cat([x,skip_connection[-3]],dim=1)

    x = self.up3(x)
    x = torch.cat([x,skip_connection[-4]],dim=1)
    
    x = self.up4(x)
    x = self.out_conv(x)
    
    return x

```