Fix up block computation for swint and convnext

Author: gitttt-1234

sleap_nn/architectures/convnext.py

@@ -5,7 +5,7 @@
 
 from functools import partial
 from typing import Any, Callable, List, Optional, Dict, Tuple
-
+import numpy as np
 import torch
 from torch import nn, Tensor
 from torchvision.ops.misc import Conv2dNormActivation
@@ -192,7 +192,7 @@ def __init__(
         else:
             self.arch = arch_types["tiny"]
 
-        self.up_blocks = len(self.arch["channels"]) - 1
+        self.up_blocks = np.log2(self.max_stride / output_stride).astype(int) - 1
         self.convs_per_block = convs_per_block
         self.stem_patch_kernel = stem_patch_kernel
         self.stem_patch_stride = stem_patch_stride

sleap_nn/architectures/swint.py

@@ -5,7 +5,7 @@
 
 from functools import partial
 from typing import Any, Callable, List, Optional, Dict, Tuple
-
+import numpy as np
 import torch
 from torch import nn
 from sleap_nn.architectures.encoder_decoder import Decoder
@@ -227,7 +227,7 @@ def __init__(
         else:
             self.arch = arch_types["tiny"]
 
-        self.up_blocks = len(self.arch["depths"]) - 1
+        self.up_blocks = np.log2(self.max_stride / output_stride).astype(int) - 1
         self.convs_per_block = convs_per_block
         self.stem_patch_stride = stem_patch_stride
         self.down_blocks = len(self.arch["depths"]) - 1