tf.squeeze

3 files changed, 1 insertions, 1 deletions

diff --git a/__pycache__/hyperparameters.cpython-38.pyc b/__pycache__/hyperparameters.cpython-38.pyc
index e42d48d1..5fc87730 100644
--- a/__pycache__/hyperparameters.cpython-38.pyc
+++ b/__pycache__/hyperparameters.cpython-38.pyc
diff --git a/__pycache__/losses.cpython-38.pyc b/__pycache__/losses.cpython-38.pyc
index 07bf7a33..703088c9 100644
--- a/__pycache__/losses.cpython-38.pyc
+++ b/__pycache__/losses.cpython-38.pyc
diff --git a/main.py b/main.py
index 87f36013..867c3f8d 100644
--- a/main.py
+++ b/main.py
@@ -41,7 +41,7 @@ def train(model):
     for i in range(hp.num_epochs):
         if i % 100 == 0:
             fn = f"checkpoint-images/img-epoch{i}"
-            save_image(fn, model.x)
+            save_image(fn, tf.squeeze(model.x))
             
         print('batch', i)
         model.train_step()