diff --git a/face_swapper/src/networks/unet.py b/face_swapper/src/networks/unet.py
index 7a6ebd6..bf8550e 100644
--- a/face_swapper/src/networks/unet.py
+++ b/face_swapper/src/networks/unet.py
@@ -25,21 +25,29 @@ class UNet(nn.Module):
 			DownSample(1024, 1024)
 		])
 
-		if self.output_size in [ 384, 512, 768, 1024 ]:
-			down_samples.append(DownSample(1024, 2048))
-		if self.output_size in [ 512, 768, 1024 ]:
-			down_samples.append(DownSample(2048, 4096))
-		if self.output_size in [ 768, 1024 ]:
-			down_samples.append(DownSample(4096, 8192))
-		if self.output_size == 1024:
-			down_samples.append(DownSample(8192, 16384))
+		if self.output_size == 512:
+			down_samples.append(DownSample(2048, 2048))
 
 		return down_samples
 
 	def create_up_samples(self) -> nn.ModuleList:
-		up_samples = nn.ModuleList(
+		up_samples = nn.ModuleList()
+
+		if self.output_size == 256:
+			up_samples.extend(
+			[
+				UpSample(1024, 1024)
+			])
+
+		if self.output_size == 512:
+			up_samples.extend(
+			[
+				UpSample(2048, 2048),
+				UpSample(4096, 1024)
+			])
+
+		up_samples.extend(
 		[
-			UpSample(1024, 1024),
 			UpSample(2048, 512),
 			UpSample(1024, 256),
 			UpSample(512, 128),
@@ -47,15 +55,6 @@ class UNet(nn.Module):
 			UpSample(128, 32)
 		])
 
-		if self.output_size in [ 384, 512, 768, 1024 ]:
-			up_samples.append(UpSample(32, 16))
-		if self.output_size in [ 512, 768, 1024 ]:
-			up_samples.append(UpSample(16, 8))
-		if self.output_size in [ 768, 1024 ]:
-			up_samples.append(UpSample(8, 4))
-		if self.output_size == 1024:
-			up_samples.append(UpSample(4, 2))
-
 		return up_samples
 
 	def forward(self, target_tensor : Tensor) -> Tuple[Tensor, ...]:
diff --git a/face_swapper/tests/test_networks.py b/face_swapper/tests/test_networks.py
index 37ccc7a..b9da7b2 100644
--- a/face_swapper/tests/test_networks.py
+++ b/face_swapper/tests/test_networks.py
@@ -5,9 +5,10 @@ from face_swapper.src.networks.aad import AAD
 from face_swapper.src.networks.unet import UNet
 
 
-@pytest.mark.parametrize('output_size', [ 256 ])
+@pytest.mark.parametrize('output_size', [ 256, 512 ])
 def test_aad_with_unet(output_size : int) -> None:
 	identity_channels = 512
+	output_channels = 2048
 	if output_size == 256:
 		output_channels = 4096
 	if output_size == 512: