diff --git a/zoobot/tensorflow/estimators/maxvit_zoobot.py b/zoobot/tensorflow/estimators/maxvit_zoobot.py
index 3f56d036..d2e4b17a 100644
--- a/zoobot/tensorflow/estimators/maxvit_zoobot.py
+++ b/zoobot/tensorflow/estimators/maxvit_zoobot.py
@@ -123,7 +123,7 @@ class Attention(tf.keras.layers.Layer):
   def __init__(self,
                hidden_size,
                head_size,
-               num_heads=None,
+               num_heads=0,
                dropatt=0.0,
                attn_axis=0,
                rel_attn_type=None,