1 files changed, 4 insertions, 0 deletions
diff --git a/cgan.py b/cgan.py
index eb0dbc3..f3bdbe5 100644
--- a/cgan.py
+++ b/cgan.py
@@ -206,6 +206,10 @@ class CGAN():
         labels_test[i*1000:] = i
         labels_val[i*500:] = i
 
+      labels_train = keras.utils.to_categorical(labels_train, 10)
+      labels_test = keras.utils.to_categorical(labels_test, 10)
+      labels_val = keras.utils.to_categorical(labels_val, 10)
+
       train_data = self.generator.predict([noise_train, labels_train])
       test_data = self.generator.predict([noise_test, labels_test])
       val_data = self.generator.predict([noise_val, labels_val])