---
                        tags:
                        - unet
                        - pix2pix
                        library_name: pytorch
                        ---
                        
                        # Pix2Pix UNet Model
                        
                        ## Model Description
                        Custom UNet model for Pix2Pix image translation.
                        - Image Size: 256
                        - Model Type: Small (256)
                        
                        ## Usage
                        
                        ```python
                        import torch
                        from small_256_model import UNet as small_UNet
                        from big_1024_model import UNet as big_UNet
                        
                        # Load the model
                        checkpoint = torch.load('model_weights.pth')
                        model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
                        model.load_state_dict(checkpoint['model_state_dict'])
                        model.eval()
                        Model Architecture
                        UNet(
  (encoder): Sequential(
    (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (5): ReLU(inplace=True)
  )
  (decoder): Sequential(
    (0): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (5): Tanh()
  )
)