broccoli-ml 0.5.4__tar.gz → 0.5.5__tar.gz

{broccoli_ml-0.5.4 → broccoli_ml-0.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 0.5.4
+Version: 0.5.5
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-0.5.4 → broccoli_ml-0.5.5}/broccoli/vit.py

@@ -215,6 +215,7 @@ class ViTEncoder(nn.Module):
                     ),  # for transformer
                 ]
             )
+            pooling_out_channels = transformer_embedding_size
 
         elif pooling_type == "max":
             self.pool = nn.Sequential(
@@ -229,6 +230,7 @@ class ViTEncoder(nn.Module):
                     ),  # for transformer
                 ]
             )
+            pooling_out_channels = transformer_embedding_size
 
         elif pooling_type == "average":
             self.pool = nn.Sequential(
@@ -243,6 +245,7 @@ class ViTEncoder(nn.Module):
                     ),  # for transformer
                 ]
             )
+            pooling_out_channels = transformer_embedding_size
 
         elif pooling_type == "concat":
 
@@ -253,9 +256,7 @@ class ViTEncoder(nn.Module):
             else:
                 self.concatpool_activation = transformer_activation()
 
-            concatpool_out_channels = (
-                pooling_kernel_voxels * cnn_activation_out_channels
-            )
+            pooling_out_channels = pooling_kernel_voxels * cnn_activation_out_channels
 
             self.pool = nn.Sequential(
                 *[
@@ -268,9 +269,6 @@ class ViTEncoder(nn.Module):
                     Rearrange(  # for transformer
                         f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
                     ),
-                    PadTensor(
-                        (0, transformer_embedding_size - concatpool_out_channels)
-                    ),
                 ]
             )
 
@@ -303,7 +301,7 @@ class ViTEncoder(nn.Module):
                 self.pool,
                 (
                     FeedforwardLayer(
-                        transformer_embedding_size,
+                        pooling_out_channels,
                         transformer_mlp_ratio,
                         transformer_embedding_size,
                         activation=transformer_activation,

{broccoli_ml-0.5.4 → broccoli_ml-0.5.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "0.5.4"
+version = "0.5.5"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}