diff --git a/csrc/codegen.cpp b/csrc/codegen.cpp
index 0bfb72f1a7a..727d25beb53 100644
--- a/csrc/codegen.cpp
+++ b/csrc/codegen.cpp
@@ -3019,14 +3019,15 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
     }
     if (loop->isUnrolled()) {
       indent() << "#pragma unroll\n";
-    } else if (
-        loop->circularBufferLoopStage() == CircularBufferLoopStage::Main) {
-      indent() << "#pragma unroll " << loop->circularBufferLoopStageDepth()
-               << "\n";
     } else if (
         loop->circularBufferLoopStage() == CircularBufferLoopStage::Epilog) {
       indent() << "#pragma unroll " << loop->circularBufferLoopStageDepth() - 1
                << "\n";
+    } else if (
+        loop->circularBufferLoopStage() !=
+        CircularBufferLoopStage::NotApplicable) {
+      indent() << "#pragma unroll " << loop->circularBufferLoopStageDepth()
+               << "\n";
     } else {
       indent() << "#pragma unroll 1\n";
     }