| 276 | } |
| 277 | |
| 278 | TensorShapeProto VecToShapeProto(const std::vector<int64>& shape_vec) { |
| 279 | TensorShapeProto shape_pb; |
| 280 | if (shape_vec.empty()) { |
| 281 | shape_pb.set_unknown_rank(true); |
| 282 | return shape_pb; |
| 283 | } |
| 284 | for (const int64 s : shape_vec) { |
| 285 | shape_pb.add_dim()->set_size(s); |
| 286 | } |
| 287 | return shape_pb; |
| 288 | } |
| 289 | |
| 290 | bool IsPlacedOnAccelerator(const string& device) { |
| 291 | return device.find("gpu") != device.npos || |