Setup KernelFallbackRequestState in c_api_tfrt in order to use KernelFallback in TF-TFRT integration.
PiperOrigin-RevId: 343595096 Change-Id: If828f781337824bd4118963ea3206a846cb8b948
This commit is contained in:
parent
03ab8d341c
commit
5d8eae5093
@ -94,14 +94,16 @@ class ScopedStepContainer {
|
||||
// prefix: optional string prefix to disambiguate step containers.
|
||||
ScopedStepContainer(const int64 step_id,
|
||||
std::function<void(const string&)> cleanup)
|
||||
: container_(strings::StrCat("__per_step_", step_id)),
|
||||
: step_id_(step_id),
|
||||
container_(strings::StrCat("__per_step_", step_id)),
|
||||
cleanup_(cleanup),
|
||||
dirty_(false) {}
|
||||
|
||||
ScopedStepContainer(const int64 step_id,
|
||||
std::function<void(const string&)> cleanup,
|
||||
const std::string& prefix)
|
||||
: container_(strings::StrCat("__", prefix, "_per_step_", step_id)),
|
||||
: step_id_(step_id),
|
||||
container_(strings::StrCat("__", prefix, "_per_step_", step_id)),
|
||||
cleanup_(cleanup),
|
||||
dirty_(false) {}
|
||||
|
||||
@ -141,8 +143,10 @@ class ScopedStepContainer {
|
||||
template <typename T>
|
||||
Status LookupOrCreate(ResourceMgr* rm, const std::string& name, T** resource,
|
||||
std::function<Status(T**)> creator) TF_MUST_USE_RESULT;
|
||||
int64 StepId() const { return step_id_; }
|
||||
|
||||
private:
|
||||
const int64 step_id_;
|
||||
const std::string container_;
|
||||
const std::function<void(const string&)> cleanup_;
|
||||
mutex mu_;
|
||||
|
Loading…
x
Reference in New Issue
Block a user