diff --git a/pysb/export/pysb_flat.py b/pysb/export/pysb_flat.py
index 0ad2a3932..56e16c0c9 100644
--- a/pysb/export/pysb_flat.py
+++ b/pysb/export/pysb_flat.py
@@ -63,7 +63,7 @@ def write_cset(cset):
         output.write("\n")
         output.write("from pysb import Model, Monomer, Parameter, Expression, "
                      "Compartment, Rule, Observable, Initial, MatchOnce, "
-                     "Annotation, ANY, WILD\n")
+                     "Annotation, MultiState, ANY, WILD\n")
         output.write("\n")
         output.write("Model()\n")
         output.write("\n")