Mark the ExpandDimOp and the SqueezeOp as non expensive
Change: 128892723
This commit is contained in:
parent
ac281d1df4
commit
6fa593f46d
@ -253,6 +253,8 @@ class ExpandDimsOp : public OpKernel {
|
||||
" and output shape ", output_shape.DebugString()));
|
||||
}
|
||||
}
|
||||
|
||||
bool IsExpensive() override { return false; }
|
||||
};
|
||||
REGISTER_KERNEL_BUILDER(Name("ExpandDims").Device(DEVICE_CPU).HostMemory("dim"),
|
||||
ExpandDimsOp);
|
||||
@ -342,6 +344,8 @@ class SqueezeOp : public OpKernel {
|
||||
}
|
||||
}
|
||||
|
||||
bool IsExpensive() override { return false; }
|
||||
|
||||
private:
|
||||
std::unordered_set<int32> squeeze_dims_;
|
||||
};
|
||||
|
Loading…
x
Reference in New Issue
Block a user