Skip to content

Commit 9e1ec8c

Browse files
authored
Enable device switching automatically for serveral operators (#8684)
1 parent ae2026e commit 9e1ec8c

File tree

3 files changed

+17
-1
lines changed

3 files changed

+17
-1
lines changed

paddle/fluid/operators/bipartite_match_op.cc

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,14 @@ class BipartiteMatchOp : public framework::OperatorWithKernel {
4141
ctx->SetOutputDim("ColToRowMatchIndices", dims);
4242
ctx->SetOutputDim("ColToRowMatchDist", dims);
4343
}
44+
45+
protected:
46+
framework::OpKernelType GetExpectedKernelType(
47+
const framework::ExecutionContext& ctx) const override {
48+
return framework::OpKernelType(
49+
framework::ToDataType(ctx.Input<LoDTensor>("DistMat")->type()),
50+
platform::CPUPlace());
51+
}
4452
};
4553

4654
template <typename T>

paddle/fluid/operators/multiclass_nms_op.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ class MultiClassNMSOp : public framework::OperatorWithKernel {
6262
return framework::OpKernelType(
6363
framework::ToDataType(
6464
ctx.Input<framework::LoDTensor>("Scores")->type()),
65-
ctx.device_context());
65+
platform::CPUPlace());
6666
}
6767
};
6868

paddle/fluid/operators/prior_box_op.cc

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,14 @@ class PriorBoxOp : public framework::OperatorWithKernel {
6767
ctx->SetOutputDim("Boxes", framework::make_ddim(dim_vec));
6868
ctx->SetOutputDim("Variances", framework::make_ddim(dim_vec));
6969
}
70+
71+
protected:
72+
framework::OpKernelType GetExpectedKernelType(
73+
const framework::ExecutionContext& ctx) const override {
74+
return framework::OpKernelType(
75+
framework::ToDataType(ctx.Input<framework::Tensor>("Input")->type()),
76+
platform::CPUPlace());
77+
}
7078
};
7179

7280
class PriorBoxOpMaker : public framework::OpProtoAndCheckerMaker {

0 commit comments

Comments
 (0)