diff --git a/src/pet/__init__.py b/src/pet/__init__.py
index 3e55142..6f832dd 100644
--- a/src/pet/__init__.py
+++ b/src/pet/__init__.py
@@ -7,6 +7,8 @@ __version__ = "0.1.0.dev0"
 from .pet_model import (
     PETModel,
     PETModelForSequenceClassification,
+    PETModelForCausalLM,
+    PETModelForSeq2SeqLM,
     PromptEncoderType,
 )
 from .tuners import (