From 81e255bfc4c3542d3c6fe182131257b42d90ec3c Mon Sep 17 00:00:00 2001
From: danleifeng <danleifeng@163.com>
Date: Thu, 6 Feb 2020 13:50:50 +0000
Subject: [PATCH] eidt image shape for fp16

---
 plsc/entry.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/plsc/entry.py b/plsc/entry.py
index a719467..12ad168 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:
-- 
GitLab