From 6415b36bcc7ab2971bc86bdb134d49be6c39e64e Mon Sep 17 00:00:00 2001
From: bolunz <chamberlain0w0@gmail.com>
Date: Fri, 22 May 2026 15:08:27 +0800
Subject: [PATCH] fix: fix cast kernel count

---
 infini_train/include/dispatcher.h | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/infini_train/include/dispatcher.h b/infini_train/include/dispatcher.h
index 29d11b73..638df76a 100644
--- a/infini_train/include/dispatcher.h
+++ b/infini_train/include/dispatcher.h
@@ -74,6 +74,9 @@ class Dispatcher {
     template <typename RetT, class... ArgsT> RetT Call(KeyT key, ArgsT... args) const {
         auto kernel = this->GetKernel(key);
         tls_autocast_context.Autocast(key, args...);
+#ifdef PROFILE_MODE
+        SetProfileContext(key.second, key.first);
+#endif
         return kernel.Call<RetT>(std::forward<ArgsT>(args)...);
     }