Fix resnet

curegit · Jun 1, 2024 · c32ccef · c32ccef
1 parent 0770ff2
commit c32ccef
Showing 1 changed file with 4 additions and 4 deletions.
diff --git 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