From 7cb9cea85367f6e1033c594e6eb280e1107cd652 Mon Sep 17 00:00:00 2001
From: suweiyue <suweiyue@baidu.com>
Date: Tue, 19 May 2020 18:21:57 +0800
Subject: [PATCH] fix neg_type bug

---
 examples/erniesage/config/erniesage_v2_cpu.yaml | 2 +-
 examples/erniesage/config/erniesage_v2_gpu.yaml | 2 +-
 examples/erniesage/dataset/graph_reader.py      | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/examples/erniesage/config/erniesage_v2_cpu.yaml b/examples/erniesage/config/erniesage_v2_cpu.yaml
index 4bdb57f..77b7805 100644
--- a/examples/erniesage/config/erniesage_v2_cpu.yaml
+++ b/examples/erniesage/config/erniesage_v2_cpu.yaml
@@ -31,7 +31,7 @@ final_fc: true
 final_l2_norm: true
 loss_type: "hinge"
 margin: 0.3
-neg_type: "random_neg"
+neg_type: "batch_neg"
 
 # infer config ------
 infer_model: "./output/last"
diff --git a/examples/erniesage/config/erniesage_v2_gpu.yaml b/examples/erniesage/config/erniesage_v2_gpu.yaml
index 4706677..7a9b4af 100644
--- a/examples/erniesage/config/erniesage_v2_gpu.yaml
+++ b/examples/erniesage/config/erniesage_v2_gpu.yaml
@@ -31,7 +31,7 @@ final_fc: true
 final_l2_norm: true
 loss_type: "hinge"
 margin: 0.3
-neg_type: "random_neg"
+neg_type: "batch_neg"
 
 # infer config ------
 infer_model: "./output/last"
diff --git a/examples/erniesage/dataset/graph_reader.py b/examples/erniesage/dataset/graph_reader.py
index 1cf2053..c811b56 100644
--- a/examples/erniesage/dataset/graph_reader.py
+++ b/examples/erniesage/dataset/graph_reader.py
@@ -73,7 +73,7 @@ class GraphGenerator(BaseDataGenerator):
         batch_src = np.array(batch_src, dtype="int64")
         batch_dst = np.array(batch_dst, dtype="int64")
 
-        if neg_type == "batch_neg":
+        if self.neg_type == "batch_neg":
             neg_shape = [1]
         else:
             neg_shape = batch_dst.shape
-- 
GitLab