| 86 | } |
| 87 | |
| 88 | void CheckpointReader::GetTensor( |
| 89 | const string& name, std::unique_ptr<tensorflow::Tensor>* out_tensor, |
| 90 | TF_Status* out_status) const { |
| 91 | Status status; |
| 92 | if (reader_ != nullptr) { |
| 93 | status = reader_->GetTensor(name, out_tensor); |
| 94 | } else { |
| 95 | tensorflow::DataType dtype; |
| 96 | tensorflow::TensorShape shape; |
| 97 | status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape); |
| 98 | if (status.ok()) { |
| 99 | out_tensor->reset(new Tensor(dtype, shape)); |
| 100 | status = v2_reader_->Lookup(name, out_tensor->get()); |
| 101 | if (!status.ok()) out_tensor->reset(); |
| 102 | } |
| 103 | } |
| 104 | if (!status.ok()) { |
| 105 | Set_TF_Status_from_Status(out_status, status); |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | std::pair<std::unique_ptr<TensorSliceReader::VarToShapeMap>, |
| 110 | std::unique_ptr<TensorSliceReader::VarToDataTypeMap>> |
no test coverage detected