change where we set the style

app.py

@@ -430,7 +430,8 @@ A tool for exploring CLIP embedding space.
         with gr.Column(scale=3, min_width=200):
             submit = gr.Button("Search embedding space")
     with gr.Row():
-        output = gr.Gallery(label="Closest images in Laion 5b using kNN", show_label=True)
+        output = gr.Gallery(label="Closest images in Laion 5b using kNN", show_label=True)\
+            .style(grid=[4,4], height="auto")
 
     embedding_base64s_state = gr.State(value=[None for i in range(max_tabs)])
     embedding_power_state = gr.State(value=[1. for i in range(max_tabs)])
@@ -483,9 +484,6 @@ A tool for exploring CLIP embedding space.
 
     # submit.click(main, inputs= [embedding_base64s[0], scale, n_samples, steps, seed], outputs=output)
     submit.click(main, inputs= [average_embedding_base64, n_samples], outputs=output)
-    # output.style(grid=[4,4], height="auto", container=True)
-    # output.style(grid=[4,4], container=True)
-    output.style(grid=4, height=336)
 
     with gr.Row():
         gr.Markdown(