diff --git a/src/latent_gemma/__init__.py b/src/latent_gemma/__init__.py
index b29343f..fb345d4 100644
--- a/src/latent_gemma/__init__.py
+++ b/src/latent_gemma/__init__.py
@@ -1,5 +1,4 @@
-from .main import LatentReasoningGemmaForCausalLM, LatentReasoningCustomModelForCausalLM, LatentReasoningForCausalLM, LatentReasoningGPT2
-
+from .main import LatentReasoningGemmaForCausalLM
 # Version info
 __version__ = "0.1.0"