broccoli-ml 0.1.35__tar.gz → 0.1.37__tar.gz

{broccoli_ml-0.1.35 → broccoli_ml-0.1.37}/PKG-INFO

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

{broccoli_ml-0.1.35 → broccoli_ml-0.1.37}/broccoli/vit.py

@@ -3,6 +3,7 @@ from typing import Optional
 
 from .transformer import TransformerEncoder
 from .cnn import ConvLayer, ConcatPool
+from .activation import ReLU, SquaredReLU, GELU, SwiGLU
 from einops import einsum
 from einops.layers.torch import Rearrange
 import torch.nn as nn
@@ -244,10 +245,27 @@ class CCT(nn.Module):
         image_classes=100,
         linear_module=nn.Linear,
         image_channels=3,
+        batch_norm=False,
     ):
 
         super().__init__()
 
+        if isinstance(cnn_activation, str):
+            cnn_activation = {
+                "ReLU": ReLU,
+                "SquaredReLU": SquaredReLU,
+                "GELU": GELU,
+                "SwiGLU": SwiGLU,
+            }[cnn_activation]
+
+        if isinstance(transformer_activation, str):
+            transformer_activation = {
+                "ReLU": ReLU,
+                "SquaredReLU": SquaredReLU,
+                "GELU": GELU,
+                "SwiGLU": SwiGLU,
+            }[transformer_activation]
+
         self.encoder = CCTEncoder(
             image_size=image_size,
             conv_kernel_size=conv_kernel_size,
@@ -275,6 +293,7 @@ class CCT(nn.Module):
             stochastic_depth=stochastic_depth,
             linear_module=linear_module,
             image_channels=image_channels,
+            batch_norm=batch_norm,
         )
         self.pool = SequencePool(
             transformer_embedding_size, linear_module, image_classes

{broccoli_ml-0.1.35 → broccoli_ml-0.1.37}/pyproject.toml

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