diff --git a/juno_samples/cava/src/cava.jn b/juno_samples/cava/src/cava.jn
index 359a83ed20db37117a8c5eef5b3e7cf31986e8cd..95a73f5b9c58539bc26d439a47df79dbf0e74c68 100644
--- a/juno_samples/cava/src/cava.jn
+++ b/juno_samples/cava/src/cava.jn
@@ -2,10 +2,8 @@ fn medianMatrix<a : number, rows, cols : usize>(m : a[rows, cols]) -> a {
   const n : usize = rows * cols;
 
   let tmp : a[rows * cols];
-  for  r = 0 to rows {
-    for c = 0 to cols {
-      tmp[r * cols + c] = m[r, c];
-    }
+  for i = 0 to rows * cols {
+    tmp[i] = m[i / cols, i % cols];
   }
 
   for i = 0 to n - 1 {