Update main.py

Added fix from Issue titu1994#27 (comment)
CaptainStabs · Nov 6, 2020 · 296b8cc · 296b8cc
1 parent 70b41fe
commit 296b8cc
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/main.py b/main.py
@@ -1,6 +1,8 @@
 import models
 import argparse
 import tensorflow as tf
+import keras
+keras.backend.set_image_data_format('channels_first')
 
 parser = argparse.ArgumentParser(description="Up-Scales an image using Image Super Resolution Model")
 parser.add_argument("imgpath", type=str, nargs="+", help="Path to input image")