From 0847bd96e9f8d4324fceac645966b4b4a5969a0b Mon Sep 17 00:00:00 2001 From: Minjae Kim Date: Sun, 15 Nov 2020 01:37:51 +0900 Subject: [PATCH] Fix apex ModuleNotFoundError (#119) --- pyprof/nvtx/nvmarker.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/pyprof/nvtx/nvmarker.py b/pyprof/nvtx/nvmarker.py index ec803c6..f2f3353 100644 --- a/pyprof/nvtx/nvmarker.py +++ b/pyprof/nvtx/nvmarker.py @@ -421,19 +421,19 @@ def patch_apex_c(): def patch_apex_pyt(): """Monkey-patch pytorch-implemented functions in APEX""" - - patch_apex_module("apex.amp") - patch_apex_module("apex.contrib.groupbn") - patch_apex_module("apex.contrib.multihead_attn") - patch_apex_module("apex.contrib.optimizers") - patch_apex_module("apex.contrib.sparsity") - patch_apex_module("apex.contrib.xentropy") - patch_apex_module("apex.fp16_utils") - patch_apex_module("apex.mlp") - patch_apex_module("apex.multi_tensor_apply") - patch_apex_module("apex.optimizers") - patch_apex_module("apex.parallel") - #patch_apex_module("apex.RNN") # Confirmed to be dead code. Do not patch + if importlib.util.find_spec("apex") is not None: + patch_apex_module("apex.amp") + patch_apex_module("apex.contrib.groupbn") + patch_apex_module("apex.contrib.multihead_attn") + patch_apex_module("apex.contrib.optimizers") + patch_apex_module("apex.contrib.sparsity") + patch_apex_module("apex.contrib.xentropy") + patch_apex_module("apex.fp16_utils") + patch_apex_module("apex.mlp") + patch_apex_module("apex.multi_tensor_apply") + patch_apex_module("apex.optimizers") + patch_apex_module("apex.parallel") + #patch_apex_module("apex.RNN") # Confirmed to be dead code. Do not patch def is_same_module_or_submodule(orig, incoming):