diff --git a/assets/js/main.js b/assets/js/main.js
index ba24e67..12efc18 100644
--- a/assets/js/main.js
+++ b/assets/js/main.js
@@ -116,7 +116,7 @@ window.addEventListener('load', function() {
     var lines = [];
     var foundTrue = false;
     // Function to process code and highlight blocks and lines
-    function processCode(index, highlight) {
+    function processCode(index) {
         // When all lines are already processed
         if (lines.length == 0) return;
         // Split into code and block ID
@@ -142,7 +142,7 @@ window.addEventListener('load', function() {
             }
         }
         // Some lines don't correspond to any block
-        if (temp[1] && index != -1 && ((!isCondition && foundTrue) || isCondition)) {
+        if (temp[1] && index != -1 && ((!isCondition && foundTrue) || isCondition || !/if/.test(lines[0]))) {
             // When sleep function, we get the timeout value from the function
             if (/\d/.test(code)) {
                 timeout = parseInt(code.replace(/[a-z\.()]/g, ''));