Change location of env var

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>
NVIDIA · Apr 17, 2024 · 3592977 · 3592977
1 parent aa36931
commit 3592977
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/nemo_curator/filters/code.py b/nemo_curator/filters/code.py
@@ -103,9 +103,6 @@ def keep_document(self, score):
 class TokenizerFertilityFilter(DocumentFilter):
 
     def __init__(self, path_to_tokenizer=None, min_char_to_token_ratio=2.5):
-        import os
-
-        os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
         from nemo.collections.common.tokenizers import SentencePieceTokenizer
 
         if path_to_tokenizer is None:

diff --git a/nemo_curator/modules/__init__.py b/nemo_curator/modules/__init__.py
@@ -11,6 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
+
+os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
 
 from .add_id import AddId
 from .exact_dedup import ExactDuplicates