| 32 | namespace grappler { |
| 33 | |
| 34 | Status GraphMemory::InferStatically( |
| 35 | const std::unordered_map<string, DeviceProperties>& devices) { |
| 36 | VirtualCluster cluster(devices); |
| 37 | TF_RETURN_IF_ERROR(cluster.Provision()); |
| 38 | TF_RETURN_IF_ERROR(cluster.Initialize(item_)); |
| 39 | RunMetadata metadata; |
| 40 | Status s = cluster.Run(item_, &metadata); |
| 41 | // The virtual cluster returns the RESOURCE_EXHAUSTED error when it detects |
| 42 | // that the model would run out of memory. We still get the metadata we need |
| 43 | // out of the simulation, so we just ignore this error. |
| 44 | if (!s.ok() && s.code() != error::RESOURCE_EXHAUSTED) { |
| 45 | return s; |
| 46 | } |
| 47 | InferFromTrace(metadata.step_stats()); |
| 48 | return Status::OK(); |
| 49 | } |
| 50 | |
| 51 | Status GraphMemory::InferDynamically(Cluster* cluster) { |
| 52 | if (!cluster->DetailedStatsEnabled()) { |