fix gpu fraction

7 years ago · e030741df9
parent 746a62ebe6
commit e030741df9
1 changed files with 2 additions and 0 deletions
--- a/paddle/contrib/inference/demo/simple_on_word2vec.cc
+++ b/paddle/contrib/inference/demo/simple_on_word2vec.cc
@ -72,6 +72,8 @@ void MainThreads(int num_threads, bool use_gpu) {
  NativeConfig config;
  config.model_dir = FLAGS_dirname + "word2vec.inference.model";
  config.use_gpu = use_gpu;
+  config.fraction_of_gpu_memory = 0.15;
+  config.device = 0;
  auto main_predictor =
      CreatePaddlePredictor<NativeConfig, PaddleEngineKind::kNative>(config);