diff --git a/models/rank/dien/reader.py b/models/rank/dien/reader.py index 16ff27a6d82adf7e5fd703bcb79cb8e06798540a..b79838c8fb1014c17319738cac0361ebc536c8b5 100755 --- a/models/rank/dien/reader.py +++ b/models/rank/dien/reader.py @@ -102,10 +102,10 @@ class Reader(ReaderBase): else: len_seq = len(b[i][0]) start_idx = random.randint(0, self.max_neg_item - len_seq - 1) - self.neg_candidate_item[start_idx:start_idx + len_seq + - 1] = b[i][0] - self.neg_candidate_cat[start_idx:start_idx + len_seq + - 1] = b[i][1] + self.neg_candidate_item[start_idx:start_idx + len_seq +1] = b[ + i][0] + self.neg_candidate_cat[start_idx:start_idx + len_seq +1] = b[ + i][1] for _ in range(len(b[i][0])): randindex = random.randint(0, len(self.neg_candidate_item) - 1)