1 files changed, 3 insertions, 2 deletions
diff --git a/tools/slang-unit-test/unit-test-compile-benchmark.cpp b/tools/slang-unit-test/unit-test-compile-benchmark.cpp
index e38edc6ad..1c0b86156 100644
--- a/tools/slang-unit-test/unit-test-compile-benchmark.cpp
+++ b/tools/slang-unit-test/unit-test-compile-benchmark.cpp
@@ -88,11 +88,12 @@ void main(uint3 threadIdx : SV_DispatchThreadID)
         SLANG_CHECK(module != nullptr);
 
         ComPtr<slang::IEntryPoint> entryPoint;
-        module->findAndCheckEntryPoint(
+        SlangResult res = module->findAndCheckEntryPoint(
             "main",
-            SLANG_STAGE_VERTEX,
+            SLANG_STAGE_COMPUTE,
             entryPoint.writeRef(),
             diagnosticBlob.writeRef());
+        SLANG_CHECK(res == SLANG_OK);
 
         slang::IComponentType* componentTypes[2] = {module, entryPoint.get()};
         ComPtr<slang::IComponentType> composedProgram;