diff --git a/pytext/torchscript/module.py b/pytext/torchscript/module.py
index a22a8edb3..495df6a4d 100644
--- a/pytext/torchscript/module.py
+++ b/pytext/torchscript/module.py
@@ -3,6 +3,10 @@
 from typing import Dict, List, Optional, Tuple
 
 import torch
+from multiray.textray.textray_details.ttypes import (  # noqa
+    QuantizationSchema,
+    FeatureSchema,
+)
 from pytext.torchscript.batchutils import (
     max_tokens,
     make_prediction_texts,