Update transformers.py

cerndb · Oct 17, 2017 · b0b156e · b0b156e · raviolli · Oct 17, 2017
1 parent 587be91
commit b0b156e
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/distkeras/transformers.py b/distkeras/transformers.py
@@ -283,9 +283,15 @@ def _transform(self, row):
         Only for internal use.
         """
         label = row[self.input_column]
-        vector = to_one_hot_encoded_dense(label, self.output_dimensionality)
-        new_row = new_dataframe_row(row, self.output_column, vector.tolist())
 
+        if(isinstance(label, types.ListType)):
+                vector = np.zeros((len(label), self.output_dimensionality))
+                for i in range(len(label)):
+                        vector[i] = to_one_hot_encoded_dense(label[i], self.output_dimensionality)
+        else:
+                vector = to_one_hot_encoded_dense(label, self.output_dimensionality)
+
+        new_row = new_dataframe_row(row, self.output_column, vector.tolist())
         return new_row
 
     def transform(self, dataframe):