From 3eacf8f77719e99f8200a3507cddfa380032380a Mon Sep 17 00:00:00 2001 From: Vibhu Jawa Date: Tue, 20 Aug 2024 16:44:17 -0700 Subject: [PATCH] Fix minor bug in max_mem_gb_classifier scripts.py Signed-off-by: Vibhu Jawa --- nemo_curator/utils/script_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nemo_curator/utils/script_utils.py b/nemo_curator/utils/script_utils.py index f854098d..663970ed 100644 --- a/nemo_curator/utils/script_utils.py +++ b/nemo_curator/utils/script_utils.py @@ -450,6 +450,7 @@ def add_distributed_classifier_cluster_args(self): self.add_distributed_args() self.add_arg_enable_spilling() self.add_arg_set_torch_to_use_rmm() + self.add_arg_max_mem_gb_classifier() # Set low default RMM pool size for classifier # to allow pytorch to grow its memory usage