diff --git a/plsc/entry.py b/plsc/entry.py
index a71946710a56039f2a750fa70156aecb3bd0abeb..12ad16817cbe090ce9e72df8de4dbf0fcc602f9d 100644
--- a/plsc/entry.py
+++ b/plsc/entry.py
@@ -386,6 +386,8 @@ class Entry(object):
         num_trainers = self.num_trainers
 
         image_shape = [int(m) for m in self.image_shape]
+        if self.data_format == "NHWC":
+            image_shape=[image_shape[1], image_shape[2], image_shape[0]]
         # model definition
         model = self.model
         if model is None: