From c32ccef1207b6e6f920c38d30b23b6425948f9ef Mon Sep 17 00:00:00 2001
From: curegit <37978051+curegit@users.noreply.github.com>
Date: Sat, 1 Jun 2024 23:34:32 +0900
Subject: [PATCH] Fix resnet

---
 descreen/networks/resnet.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/descreen/networks/resnet.py b/descreen/networks/resnet.py
index c46de95..ebc2e75 100644
--- a/descreen/networks/resnet.py
+++ b/descreen/networks/resnet.py
@@ -21,13 +21,13 @@ def forward(self, x: Tensor) -> Tensor:
 
     def input_size_unchecked(self, output_size: int) -> int:
         size = output_size
-        for _ in range(len(self.blocks)):
-            size = input_size(size, 3)
+        for b in self.blocks:
+            size = b.input_size(size)
         return size
 
     def output_size_unchecked(self, input_size: int) -> int:
         size = input_size
-        for _ in range(len(self.blocks)):
-            size = output_size(size, 3)
+        for b in self.blocks:
+            size = b.output_size(size)
         return size