diff --git a/jraph/__init__.py b/jraph/__init__.py
index 1bc1571..836ca1b 100644
--- a/jraph/__init__.py
+++ b/jraph/__init__.py
@@ -71,7 +71,7 @@
 from jraph._src.utils import zero_out_padding
 
 
-__version__ = "0.0.5.dev0"
+__version__ = "0.0.6.dev0"
 
 __all__ = ("ArrayTree", "DeepSets", "GraphConvolution", "GraphMapFeatures",
            "InteractionNetwork", "RelationNetwork", "GraphNetGAT", "GAT",