File tree Expand file tree Collapse file tree 1 file changed +5
-3
lines changed
paddle/fluid/operators/distributed_ops Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Original file line number Diff line number Diff line change @@ -31,14 +31,16 @@ class SplitByrefOp : public framework::OperatorWithKernel {
31
31
auto in_dims = ctx->GetInputDim (" X" );
32
32
auto outs_names = ctx->Outputs (" Out" );
33
33
size_t num = static_cast <size_t >(ctx->Attrs ().Get <int >(" num" ));
34
- std::vector<int > sections = static_cast <std::vector<int >>(
35
- ctx->Attrs ().Get <std::vector<int >>(" sections" ));
34
+ auto sections = ctx->Attrs ().Get <std::vector<int >>(" sections" );
36
35
const size_t outs_number = outs_names.size ();
37
36
std::vector<framework::DDim> outs_dims;
38
37
outs_dims.reserve (outs_number);
39
38
40
39
if (num > 0 ) {
41
- int64_t in_axis_dim = in_dims[0 ];
40
+ int64_t in_axis_dim = 0 ;
41
+ if (ctx->IsRuntime ()) {
42
+ in_axis_dim = in_dims[0 ];
43
+ }
42
44
PADDLE_ENFORCE_EQ (in_axis_dim % num, 0 ,
43
45
" tensor split does not result"
44
46
" in an equal division" );
You can’t perform that action at this time.
0 commit comments