broccoli-ml 1.2.1__tar.gz → 1.2.2__tar.gz

{broccoli_ml-1.2.1 → broccoli_ml-1.2.2}/PKG-INFO

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

{broccoli_ml-1.2.1 → broccoli_ml-1.2.2}/broccoli/vit.py

@@ -90,8 +90,22 @@ class SequencePoolClassificationHead(ClassificationHead):
         as a generalisation of average pooling.
     """
 
-    def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
-        super().__init__(d_model, linear_module, out_dim, batch_norm_logits=batch_norm)
+    def __init__(
+        self,
+        d_model,
+        n_classes,
+        linear_module=nn.Linear,
+        logit_projection_layer=nn.Linear,
+        batch_norm_logits=True,
+    ):
+        super().__init__(
+            d_model,
+            n_classes,
+            linear_module=linear_module,
+            logit_projection_layer=logit_projection_layer,
+            batch_norm_logits=batch_norm_logits,
+        )
+
         self.summarize = SequencePool(d_model, linear_module)
         # Rebuild the classification process with the correct summary module:
         self.classification_process = nn.Sequential(

{broccoli_ml-1.2.1 → broccoli_ml-1.2.2}/pyproject.toml

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