diff --git a/src/llm_benchmark/control/double.py b/src/llm_benchmark/control/double.py
index 4be41d7..bb5caf3 100644
--- a/src/llm_benchmark/control/double.py
+++ b/src/llm_benchmark/control/double.py
@@ -4,20 +4,15 @@
 class DoubleForLoop:
     @staticmethod
     def sum_square(n: int) -> int:
-        """Sum of squares of numbers from 0 to n (exclusive)
-
-        Args:
-            n (int): Number to sum up to
-
-        Returns:
-            int: Sum of squares of numbers from 0 to n
-        """
-        sum_ = 0
-        for i in range(n):
-            for j in range(n):
-                if i == j:
-                    sum_ += i * j
-        return sum_
+            """Sum of squares of numbers from 0 to n (exclusive)
+    
+            Args:
+                n (int): Number to sum up to
+    
+            Returns:
+                int: Sum of squares of numbers from 0 to n
+            """
+            return sum(i * i for i in range(n))
 
     @staticmethod
     def sum_triangle(n: int) -> int:
@@ -90,4 +85,4 @@ def sum_matrix(m: List[List[int]]) -> int:
         for i in range(len(m)):
             for j in range(len(m[i])):
                 sum_ += m[i][j]
-        return sum_
+        return sum_
\ No newline at end of file