diff --git a/vizro-ai/examples/dashboard_ui/app.py b/vizro-ai/examples/dashboard_ui/app.py
index 38dc6099d..aa4f27a2b 100644
--- a/vizro-ai/examples/dashboard_ui/app.py
+++ b/vizro-ai/examples/dashboard_ui/app.py
@@ -15,6 +15,7 @@
     CodeClipboard,
     CustomDashboard,
     DropdownMenu,
+    FlexContainer,
     HeaderComponent,
     Icon,
     Modal,
@@ -24,7 +25,6 @@
     UserPromptTextArea,
     UserUpload,
     custom_table,
-    FlexContainer
 )
 from dash import Input, Output, State, callback, ctx, dcc, get_asset_url, html
 from vizro import Vizro
@@ -127,8 +127,9 @@
                     ),
                     components=[
                         UserPromptTextArea(id="text-area"),
-                        MyDropdown(options=SUPPORTED_MODELS["OpenAI"], value="gpt-4o-mini", multi=False,
-                                   id="model-dropdown"),
+                        MyDropdown(
+                            options=SUPPORTED_MODELS["OpenAI"], value="gpt-4o-mini", multi=False, id="model-dropdown"
+                        ),
                         vm.Button(
                             id="trigger-button",
                             text="Run Vizro-AI",
diff --git a/vizro-ai/examples/dashboard_ui/assets/custom_css.css b/vizro-ai/examples/dashboard_ui/assets/custom_css.css
index c21fd7406..0eb9a42a0 100644
--- a/vizro-ai/examples/dashboard_ui/assets/custom_css.css
+++ b/vizro-ai/examples/dashboard_ui/assets/custom_css.css
@@ -258,6 +258,7 @@
 
 #dropdown-menu-div {
   align-items: center;
+  align-self: flex-end;
   border: 0.5px solid gray;
   border-radius: 8px;
   display: flex;
@@ -266,7 +267,6 @@
   min-width: 130px;
   padding-left: 4px;
   width: 130px;
-  align-self: flex-end;
 }
 
 #custom-header-div {
@@ -343,10 +343,9 @@ img#banner {
 
 .flex-container {
   display: flex;
-  flex-wrap: wrap;
+  flex-flow: column wrap;
   gap: 12px;
   justify-content: flex-start;
-  flex-direction: column;
   overflow: none;
 }
 
@@ -362,7 +361,6 @@ img#banner {
   width: 100%;
 }
 
-
 .flex-container #graph {
   height: 712px;
-}
\ No newline at end of file
+}
diff --git a/vizro-ai/examples/dashboard_ui/components.py b/vizro-ai/examples/dashboard_ui/components.py
index 34f9ed2ef..793f27dcc 100644
--- a/vizro-ai/examples/dashboard_ui/components.py
+++ b/vizro-ai/examples/dashboard_ui/components.py
@@ -391,17 +391,17 @@ def build(self):
             toggleClassName="dropdown-menu-toggle-class",
         )
         download_div = html.Div(
-                        children=[
-                            html.Span("download", className="material-symbols-outlined", id=f"{self.id}-icon"),
-                            dropdown_menu,
-                            dbc.Tooltip(
-                                "Download this plot to your device as a plotly JSON or interactive HTML file "
-                                "for easy sharing or future use.",
-                                target="dropdown-menu-div",
-                            ),
-                        ],
-                        id="dropdown-menu-div",
-                    )
+            children=[
+                html.Span("download", className="material-symbols-outlined", id=f"{self.id}-icon"),
+                dropdown_menu,
+                dbc.Tooltip(
+                    "Download this plot to your device as a plotly JSON or interactive HTML file "
+                    "for easy sharing or future use.",
+                    target="dropdown-menu-div",
+                ),
+            ],
+            id="dropdown-menu-div",
+        )
         return download_div