From 786adf73a2acfd0c7d63ba7a584810b281028c49 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 5 Mar 2025 10:27:55 +0100 Subject: [PATCH] Fix UnetPro --- face_swapper/src/networks/unet.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/face_swapper/src/networks/unet.py b/face_swapper/src/networks/unet.py index 2f185bd..7a6ebd6 100644 --- a/face_swapper/src/networks/unet.py +++ b/face_swapper/src/networks/unet.py @@ -80,8 +80,8 @@ class UNet(nn.Module): class UNetPro(UNet): - def __init__(self) -> None: - super(UNet, self).__init__() + def __init__(self, output_size : int) -> None: + super().__init__(output_size) self.resnet = models.resnet34(weights = ResNet34_Weights.DEFAULT) self.down_samples = self.create_down_samples() self.up_samples = self.create_up_samples()