broccoli-ml 4.0.0__tar.gz → 4.0.1__tar.gz

{broccoli_ml-4.0.0 → broccoli_ml-4.0.1}/PKG-INFO

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

{broccoli_ml-4.0.0 → broccoli_ml-4.0.1}/broccoli/transformer.py

@@ -233,7 +233,7 @@ class MHAttention(nn.Module):
                 q,
                 k,
                 v,
-                dropout_p=self.dropout if self.training else 0.0,
+                dropout_p=self.dropout.p if self.training else 0.0,
                 softmax_scale=scaling_factor,
                 causal=self.causal,
             )
@@ -257,6 +257,8 @@ class MHAttention(nn.Module):
 
             qk_scores = F.softmax(qk_scores, dim=-1)
 
+            qk_scores = self.dropout(qk_scores)
+
             output_with_heads = qk_scores @ v
 
             output_without_heads = rearrange(output_with_heads, "b h t d -> b t (h d)")

{broccoli_ml-4.0.0 → broccoli_ml-4.0.1}/pyproject.toml

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