| 20 | |
| 21 | #ifdef WITH_CUDA |
| 22 | TEST(Api, embedding_test) { |
| 23 | EnvScope scope; |
| 24 | Device device("cuda"); |
| 25 | Graph graph = Graph::Load("/path/to/embedding", device); |
| 26 | int64_t batch_size = 10000; |
| 27 | int64_t num_features = 39; |
| 28 | |
| 29 | std::vector<int64_t> data(batch_size * num_features); |
| 30 | std::fill(data.begin(), data.end(), 1); |
| 31 | std::vector<Tensor> inputs; |
| 32 | inputs.emplace_back( |
| 33 | Tensor::from_buffer(data.data(), Shape({batch_size, num_features}), device, DType::kInt64)); |
| 34 | |
| 35 | const auto& value = graph.Forward(inputs); |
| 36 | |
| 37 | ASSERT_TRUE(value.IsTensor()); |
| 38 | Tensor output = value.ToTensor(); |
| 39 | Shape shape = output.shape(); |
| 40 | ASSERT_EQ(shape.At(0), batch_size); |
| 41 | ASSERT_EQ(shape.At(1), 1); |
| 42 | |
| 43 | std::vector<float> buf(batch_size); |
| 44 | output.copy_to(buf.data()); |
| 45 | } |
| 46 | #endif |
| 47 | |
| 48 | } // namespace oneflow_api |