Update iterator.py

apache · Jan 19, 2017 · 571b9a4 · 571b9a4
1 parent eaae080
commit 571b9a4
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/example/ssd/dataset/iterator.py b/example/ssd/dataset/iterator.py
@@ -45,12 +45,7 @@ def __init__(self, imdb, batch_size, data_shape, \
         if isinstance(data_shape, int):
             data_shape = (data_shape, data_shape)
         self._data_shape = data_shape
-        if isinstance(mean_pixels, list):
-            self._mean_pixels = mx.nd.Reshape(
-                    mx.nd.array(mean_pixels),
-                    shape = (3,1,1))
-        else:
-            self._mean_pixels = mean_pixels
+        self._mean_pixels = mx.nd.array(mean_pixels).reshape((3,1,1))
         if not rand_samplers:
             self._rand_samplers = []
         else: