Disable nvFuser setup with PyTorch 23.11 and later (NVIDIA#9837) (NVI…

…DIA#9870) * Disable nvFuser setup with PyTorch 23.11 and later * Apply isort and black reformatting --------- Signed-off-by: Abhishree <[email protected]> Signed-off-by: athitten <[email protected]> Co-authored-by: Abhishree Thittenamane <[email protected]> Co-authored-by: athitten <[email protected]>
Dido0o0 · Aug 23, 2024 · 06f1604 · 06f1604
1 parent 8f3ef8c
commit 06f1604
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_base_model.py b/nemo/collections/nlp/models/language_modeling/megatron_base_model.py
@@ -408,7 +408,9 @@ def is_official_release_version(nvidia_torch_version):
                 self.cfg.persist_layer_norm = False
 
             # NVFUSER available starting with 21.11
-            if NVIDIA_TORCH_MAJOR >= 21 or (NVIDIA_TORCH_MAJOR == 21 and NVIDIA_TORCH_MINOR >= 11):
+            if (NVIDIA_TORCH_MAJOR >= 21 or (NVIDIA_TORCH_MAJOR == 21 and NVIDIA_TORCH_MINOR >= 11)) and (
+                NVIDIA_TORCH_MAJOR < 23 or (NVIDIA_TORCH_MAJOR == 23 and NVIDIA_TORCH_MINOR < 11)
+            ):
 
                 # NVFUSER
                 torch._C._jit_set_profiling_executor(True)