Monkey patch functions in APEX.

Signed-off-by: Aditya Agrawal <[email protected]>

pyprof/nvtx/nvmarker.py

@@ -334,12 +334,28 @@ def new_iter(self, *args, **kwargs):
 
     mod.DataLoader.__iter__ = new_iter
 
+# Monkey-patch functions in APEX
+#
+def patch_apex():
+    import importlib
+    if importlib.util.find_spec("amp_C") is not None:
+        import amp_C
+        patchClass(amp_C)
+
+    if importlib.util.find_spec("fused_adam_cuda") is not None:
+        import fused_adam_cuda
+        patchClass(fused_adam_cuda)
+
+    if importlib.util.find_spec("fused_layer_norm_cuda") is not None:
+        import fused_layer_norm_cuda
+        patchClass(fused_layer_norm_cuda)
 
 def init():
     print("Initializing NVTX monkey patches")
 
     patch_dataloader()
     patch_torch_classes()
     patch_torch_nn_forward_functions()
+    patch_apex()
 
     print("Done with NVTX monkey patching")