diff --git a/pytext/data/tensorizers.py b/pytext/data/tensorizers.py
index fb134691b..07b4eae07 100644
--- a/pytext/data/tensorizers.py
+++ b/pytext/data/tensorizers.py
@@ -308,12 +308,13 @@ def tensorize(self, batch):
         return pad_and_tensorize(batch, dtype=torch.float)
 
 
-class JsonFloatListTensorizer(Tensorizer):
+class FloatListTensorizer(Tensorizer):
     """Numberize numeric labels."""
 
     class Config(Tensorizer.Config):
         #: The name of the label column to parse from the data source.
         column: str
+        dim: Optional[int] = None
 
     @classmethod
     def from_config(cls, config: Config):
@@ -324,7 +325,7 @@ def __init__(self, column: str):
         self.column = column
 
     def numberize(self, row):
-        res = json.loads(row[self.column])
+        res = json.loads(row[self.column].replace(" ", ","))
         if type(res) is not list:
             raise ValueError(f"{res} is not a valid float list")
         return [float(n) for n in res]
diff --git a/pytext/models/doc_model.py b/pytext/models/doc_model.py
index 4c6ff2140..247bb3c6e 100644
--- a/pytext/models/doc_model.py
+++ b/pytext/models/doc_model.py
@@ -64,6 +64,9 @@ class ModelInput(Model.Config.ModelInput):
         inputs: ModelInput = ModelInput()
         embedding: WordEmbedding.Config = WordEmbedding.Config()
 
+    input_names = ["tokens", "tokens_lens"]
+    output_names = ["scores"]
+
     def arrange_model_inputs(self, tensor_dict):
         tokens, seq_lens = tensor_dict["tokens"]
         return (tokens, seq_lens)
@@ -71,13 +74,16 @@ def arrange_model_inputs(self, tensor_dict):
     def arrange_targets(self, tensor_dict):
         return tensor_dict["labels"]
 
+    def vocab_to_export(self, tensorizers):
+        return {"tokens": list(tensorizers["tokens"].vocab)}
+
     def caffe2_export(self, tensorizers, tensor_dict, path, export_onnx_path=None):
         exporter = ModelExporter(
             ModelExporter.Config(),
-            ["tokens", "tokens_lens"],
+            self.input_names,
             self.arrange_model_inputs(tensor_dict),
-            {"tokens": list(tensorizers["tokens"].vocab)},
-            ["scores"],
+            self.vocab_to_export(tensorizers),
+            self.output_names,
         )
         return exporter.export_to_caffe2(self, path, export_onnx_path=export_onnx_path)
 
@@ -88,6 +94,12 @@ def create_embedding(cls, config: Config, tensorizers: Dict[str, Tensorizer]):
             len(vocab), config.embedding.embed_dim, None, None, vocab.idx[UNK], []
         )
 
+    @classmethod
+    def create_decoder(cls, config: Config, representation_dim: int, num_labels: int):
+        return create_module(
+            config.decoder, in_dim=representation_dim, out_dim=num_labels
+        )
+
     @classmethod
     def from_config(cls, config: Config, tensorizers: Dict[str, Tensorizer]):
         labels = tensorizers["labels"].labels
@@ -95,10 +107,8 @@ def from_config(cls, config: Config, tensorizers: Dict[str, Tensorizer]):
         representation = create_module(
             config.representation, embed_dim=embedding.embedding_dim
         )
-        decoder = create_module(
-            config.decoder,
-            in_dim=representation.representation_dim,
-            out_dim=len(labels),
+        decoder = cls.create_decoder(
+            config, representation.representation_dim, len(labels)
         )
         # TODO change from_config function of ClassificationOutputLayer after migriting to new design
         output_layer = ClassificationOutputLayer(
diff --git a/pytext/task/new_task.py b/pytext/task/new_task.py
index 0702c2399..c01991e82 100644
--- a/pytext/task/new_task.py
+++ b/pytext/task/new_task.py
@@ -8,7 +8,6 @@
 from pytext.config.component import ComponentType, create_component, create_trainer
 from pytext.data.data import Data
 from pytext.data.tensorizers import Tensorizer
-from pytext.exporters import ModelExporter
 from pytext.metric_reporters import (
     ClassificationMetricReporter,
     MetricReporter,
@@ -179,8 +178,8 @@ def export(self, model, export_path, metric_channels=None, export_onnx_path=None
         precision.deactivate(model)
 
         batch = next(iter(self.data.batches(Stage.TRAIN)))
-        print("Saving caffe2 model to: " + export_path)
-        model.caffe2_export(
+        print(f"Saving caffe2 model to: {export_path}")
+        return model.caffe2_export(
             self.data.tensorizers, batch, export_path, export_onnx_path=export_onnx_path
         )