From ea7af61ed4d86da9acd313e3ea01fc1599fa9764 Mon Sep 17 00:00:00 2001 From: lessw2020 Date: Mon, 4 Mar 2024 19:51:35 -0800 Subject: [PATCH] update is_comm_kernel check to work with newer nccl versions --- hta/utils/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hta/utils/utils.py b/hta/utils/utils.py index 80d10f4e..f7355398 100644 --- a/hta/utils/utils.py +++ b/hta/utils/utils.py @@ -60,7 +60,7 @@ def is_comm_kernel(name: str) -> bool: Returns: A boolean indicating if the kernel is a communication kernel. """ - return "ncclKernel" in name + return "ncclKernel" in name or "ncclDev" in name def is_memory_kernel(name: str) -> bool: