Fix Noisy CUDA Shutdown (#20)

* Reorder imports Signed-off-by: Ryan Wolf <[email protected]> * Update precommit exclude Signed-off-by: Ryan Wolf <[email protected]> * Clarify comment Signed-off-by: Ryan Wolf <[email protected]> --------- Signed-off-by: Ryan Wolf <[email protected]>
NVIDIA · Apr 2, 2024 · 35d5f8f · 35d5f8f
1 parent fc167a6
commit 35d5f8f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -44,4 +44,4 @@ repos:
     hooks:
       - id: isort
         name: Format imports
-        exclude: docs/
+        exclude: docs/|nemo_curator/modules/__init__.py
diff --git a/nemo_curator/modules/__init__.py b/nemo_curator/modules/__init__.py
@@ -13,14 +13,18 @@
 # limitations under the License.
 
 from .add_id import AddId
-from .distributed_data_classifier import DomainClassifier, QualityClassifier
 from .exact_dedup import ExactDuplicates
 from .filter import Filter, Score, ScoreFilter
 from .fuzzy_dedup import LSH, MinHash
 from .meta import Sequential
 from .modify import Modify
 from .task import TaskDecontamination
 
+# Pytorch related imports must come after all imports that require cugraph,
+# because of context cleanup issues b/w pytorch and cugraph
+# See this issue: https://github.com/rapidsai/cugraph/issues/2718
+from .distributed_data_classifier import DomainClassifier, QualityClassifier
+
 __all__ = [
     "DomainClassifier",
     "ExactDuplicates",