提交 4520afcf 编写于 作者: Y yangyaming

Consider corner case.

上级 f2d596d4
......@@ -109,12 +109,24 @@ class ExpandGradKernel : public framework::OpKernel {
}
int dims = reshape_dims_vec.size() * 6 + reduce_dims_vec.size() - 7;
// no need reduce, just copy
if (reduce_dims_vec.size() == 0) {
auto* in0 = context.Input<Tensor>(framework::GradVarName("Out"));
auto* out0 = context.Output<Tensor>(framework::GradVarName("X"));
out0->mutable_data<T>(context.GetPlace());
if (platform::is_cpu_place(context.GetPlace())) {
out0->CopyFrom<T>(*in0, platform::CPUPlace());
} else {
out0->CopyFrom<T>(*in0, platform::GPUPlace());
}
} else {
switch (dims) {
REP_EXPAND_GRAD_TEMPLATE(72)
default:
PADDLE_ENFORCE(false, "Only support tensor whose rank in [1, 6].");
};
}
}
protected:
template <int Dims>
......
......@@ -22,8 +22,8 @@ class TestExpandOpRank2(OpTest):
def setUp(self):
self.op_type = "expand"
self.inputs = {'X': np.random.random((12, 14)).astype("float32")}
self.attrs = {'expandTimes': [3, 4]}
output = np.tile(self.inputs['X'], (3, 4))
self.attrs = {'expandTimes': [1, 1]}
output = np.tile(self.inputs['X'], (1, 1))
self.outputs = {'Out': output}
def test_check_output(self):
......@@ -37,8 +37,8 @@ class TestExpandOpRank3(OpTest):
def setUp(self):
self.op_type = "expand"
self.inputs = {'X': np.random.random((2, 4, 5)).astype("float32")}
self.attrs = {'expandTimes': [3, 2, 1]}
output = np.tile(self.inputs['X'], (3, 2, 1))
self.attrs = {'expandTimes': [1, 1, 1]}
output = np.tile(self.inputs['X'], (1, 1, 1))
self.outputs = {'Out': output}
def test_check_output(self):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册