fix params with only 1 dim (#15828) (#15832)

* fix params with only 1 dim * test=develop
PaddlePaddle · Feb 21, 2019 · 4b3f9e5 · 4b3f9e5
1 parent 2307baf
commit 4b3f9e5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/python/paddle/fluid/io.py b/python/paddle/fluid/io.py
@@ -766,7 +766,10 @@ def __load_persistable_vars(executor, dirname, need_load_vars):
                     dtype=slice_var.dtype,
                     persistable=True)
 
-                dim1_flatten = reduce(lambda x, y: x * y, slice.shape[1:])
+                dim1_flatten = 1
+                if len(slice.shape) >= 2:
+                    dim1_flatten = reduce(lambda x, y: x * y, slice.shape[1:])
+
                 start = int(offset / dim1_flatten)
                 end = int(offset / dim1_flatten + slice.shape[0])
 

diff --git a/python/paddle/fluid/transpiler/distribute_transpiler.py b/python/paddle/fluid/transpiler/distribute_transpiler.py
@@ -1020,7 +1020,11 @@ def _get_slice_var_info(self, slice_var):
         skip_dim0 = 0
         slice_vars = self.param_var_mapping[orig_var_name]
 
-        orig_dim1_flatten = reduce(lambda x, y: x * y, slice_vars[0].shape[1:])
+        orig_dim1_flatten = 1
+
+        if len(slice_vars[0].shape) >= 2:
+            orig_dim1_flatten = reduce(lambda x, y: x * y,
+                                       slice_vars[0].shape[1:])
 
         for slice_var in slice_vars[:block_idx]:
             skip_dim0 += slice_var.shape[0]