diff --git a/data/gan.py b/data/gan.py
index f2f3cdb..d1b3123 100644
--- a/data/gan.py
+++ b/data/gan.py
@@ -625,6 +625,7 @@ class Predict(GNet):
                 
                 tf.compat.v1.reset_default_graph()
                 df = pd.DataFrame(df)
+                df.columns = columns
                 return df.to_dict(orient='list')
                         # return df.to_dict(orient='list')
                         # count = str(len(os.listdir(self.out_dir)))