edit final comment

app.py

@@ -139,7 +139,7 @@ def check_validity(dim1,dim2):
 
 
 def substitute_ones_with_concat(dim1,dim2):
-    for i in range(len(dim1)-2):
+    for i in range(len(dim1)-1):
         dim1[i] = dim2[i] if dim1[i] == 1 else dim1[i]
         dim2[i] = dim1[i] if dim2[i] == 1 else dim2[i]
     return dim1, dim2
@@ -168,7 +168,9 @@ def predict(dim1, dim2):
     out += "* all the other dimensions should be equal to one another (green or red colors)\n\n" + table3
     if "X" not in checks : 
         dim1[-1] = dim2[-1]
-        out += f"\n# Final dimension\n `output.shape = {dim1}`"
+        out += "\n# Final dimension\n"
+        out+="as highlighted in <strong style='color:green'> green </strong> \n\n"
+        out+= f"`output.shape = {dim1}`"
     return out
 
 
@@ -176,7 +178,7 @@ demo = gr.Interface(
     predict,
     inputs=["text", "text"],
     outputs=["markdown"],
-    examples=[["9,2,1,3,3", "5,3,7"], ["1,2,3", "5,2,7"]],
+    examples=[["9,2,1,3,3", "5,3,7"], ["7,4,2,3", "5,2,7"]],
 )
 
 demo.launch(debug=True)