Fix

model.py

@@ -123,8 +123,8 @@ class Model:
         d = nrows * ncols - len(images)
         if d > 0:
             images += [np.full((h, w, 3), 255, dtype=np.uint8)] * d
-        grid = np.asarray(images).reshape(2, 2, h, w, 3).transpose(
-            0, 2, 1, 3, 4).reshape(2 * h, 2 * w, 3)
+        grid = np.asarray(images).reshape(nrows, ncols, h, w, 3).transpose(
+            0, 2, 1, 3, 4).reshape(nrows * h, ncols * w, 3)
         return PIL.Image.fromarray(grid)
 
     def run_simple(self) -> PIL.Image.Image: