Merge pull request #2117 from daviddavo/staging

Fix ImplicitCF failing if not using stratified split #2024
recommenders-team · Jun 25, 2024 · 9f630e3 · 9f630e3
2 parents 211ee70 + cb2b282
commit 9f630e3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/recommenders/models/deeprec/DataModel/ImplicitCF.py b/recommenders/models/deeprec/DataModel/ImplicitCF.py
@@ -80,6 +80,7 @@ def _data_processing(self, train, test):
             user_idx = df[[self.col_user]].drop_duplicates().reindex()
             user_idx[self.col_user + "_idx"] = np.arange(len(user_idx))
             self.n_users = len(user_idx)
+            self.n_users_in_train = train[self.col_user].nunique()
             self.user_idx = user_idx
 
             self.user2id = dict(
@@ -210,7 +211,7 @@ def sample_neg(x):
                 if neg_id not in x:
                     return neg_id
 
-        indices = range(self.n_users)
+        indices = range(self.n_users_in_train)
         if self.n_users < batch_size:
             users = [random.choice(indices) for _ in range(batch_size)]
         else: