| 206 | } |
| 207 | |
| 208 | IValue Graph::Forward(const IValue& inputs) { |
| 209 | std::vector<Tensor> input_tensors; |
| 210 | if (inputs.IsNone()) { |
| 211 | // do nothing |
| 212 | } else if (inputs.IsTensor()) { |
| 213 | input_tensors.emplace_back(inputs.ToTensor()); |
| 214 | } else if (inputs.IsTensorVector()) { |
| 215 | input_tensors = inputs.ToTensorVector(); |
| 216 | } else { |
| 217 | LOG(WARNING) << "Graph currently only support types: Tensor/vector(Tensor)/None"; |
| 218 | } |
| 219 | |
| 220 | std::vector<Tensor> output_tensors = graph_->Forward(input_tensors); |
| 221 | if (output_tensors.empty()) { |
| 222 | return IValue{}; |
| 223 | } else if (output_tensors.size() == 1) { |
| 224 | return IValue(output_tensors.at(0)); |
| 225 | } else { |
| 226 | return IValue(output_tensors); |
| 227 | } |
| 228 | } |
| 229 | |
| 230 | void Graph::set_batch_size(int batch_size) { graph_->set_batch_size(batch_size); } |
| 231 | |