From f2314cd913b4681b82b69c8a133b13660951bb9f Mon Sep 17 00:00:00 2001 From: JalenCato Date: Mon, 25 Aug 2025 21:13:14 +0000 Subject: [PATCH] add --- python/graphstorm/gconstruct/transform.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/graphstorm/gconstruct/transform.py b/python/graphstorm/gconstruct/transform.py index 69aee9c094..b3af32b4ed 100644 --- a/python/graphstorm/gconstruct/transform.py +++ b/python/graphstorm/gconstruct/transform.py @@ -43,6 +43,7 @@ LABEL_STATS_FIELD = "training_label_stats" LABEL_STATS_FREQUENCY_COUNT = "frequency_cnt" +MAX_CATEGORIES = 100 CLASSIFICATION_LABEL_STATS_TYPES = [LABEL_STATS_FREQUENCY_COUNT] @@ -532,7 +533,10 @@ def update_info(self, info): assert len(info) == 0 return - self._val_dict = {str(key): i for i, key in enumerate(np.unique(np.concatenate(info)))} + # Set the maximum categories number to avoid OOM on shared memory + self._val_dict = {str(key): min(i, MAX_CATEGORIES - 1) + for i, key in enumerate(np.unique(np.concatenate(info)))} + # We need to save the mapping in the config object. if self._conf is not None: self._conf['mapping'] = self._val_dict