| 203 | TensorSliceReader::~TensorSliceReader() { gtl::STLDeleteValues(&tensors_); } |
| 204 | |
| 205 | bool TensorSliceReader::HasTensor(const string& name, TensorShape* shape, |
| 206 | DataType* type) const { |
| 207 | mutex_lock l(mu_); |
| 208 | const TensorSliceSet* tss = gtl::FindPtrOrNull(tensors_, name); |
| 209 | if (!tss && !all_shards_loaded_) { |
| 210 | VLOG(1) << "Did not find tensor in preferred shard, loading all shards: " |
| 211 | << name; |
| 212 | LoadAllShards(); |
| 213 | tss = gtl::FindPtrOrNull(tensors_, name); |
| 214 | } |
| 215 | if (tss) { |
| 216 | if (shape) { |
| 217 | *shape = tss->shape(); |
| 218 | } |
| 219 | if (type) { |
| 220 | *type = tss->type(); |
| 221 | } |
| 222 | return true; |
| 223 | } else { |
| 224 | return false; |
| 225 | } |
| 226 | } |
| 227 | |
| 228 | Status TensorSliceReader::GetTensor( |
| 229 | const string& name, std::unique_ptr<tensorflow::Tensor>* out_tensor) const { |