diff --git a/tests/layout/test_flex.py b/tests/layout/test_flex.py
index 159edcad8..13a90c609 100644
--- a/tests/layout/test_flex.py
+++ b/tests/layout/test_flex.py
@@ -455,9 +455,9 @@ def test_flex_item_min_height():
 @assert_no_logs
 def test_flex_auto_margin():
     # Regression test for https://github.com/Kozea/WeasyPrint/issues/800
-    page, = render_pages('<div style="display: flex; margin: auto">')
+    page, = render_pages('<div id="div1" style="display: flex; margin: auto">')
     page, = render_pages(
-        '<div style="display: flex; flex-direction: column; margin: auto">')
+        '<div id="div2" style="display: flex; flex-direction: column; margin: auto">')
 
 
 @assert_no_logs
@@ -726,3 +726,31 @@ def test_flex_column_width():
     assert content.width == paper.width
     assert toppart.width == paper.width
     assert bottompart.position_y > toppart.position_y + toppart.height
+
+
+@assert_no_logs
+def test_flex_auto_margin2():
+    # Regression test for https://github.com/Kozea/WeasyPrint/issues/2054
+    page, = render_pages('''
+<style>
+  #outer {
+    background: red;
+  }
+  #inner {
+    margin: auto;
+    display: flex;
+    width: 160px;
+    height: 160px;
+    background: black;
+  }
+</style>
+
+<div id="outer">
+  <div id="inner"></div>
+</div>
+''')
+    html, = page.children
+    body, = html.children
+    outer, = body.children
+    inner, = outer.children
+    assert inner.margin_left != 0
diff --git a/weasyprint/layout/flex.py b/weasyprint/layout/flex.py
index c039c6a5b..26697e7ec 100644
--- a/weasyprint/layout/flex.py
+++ b/weasyprint/layout/flex.py
@@ -35,7 +35,6 @@ def flex_layout(context, box, bottom_space, skip_stack, containing_block,
     else:
         axis, cross = 'height', 'width'
     LOGGER.debug("%r 9.2 Line Length Determination", box)
-    # breakpoint()
 
     margin_left = 0 if box.margin_left == 'auto' else box.margin_left
     margin_right = 0 if box.margin_right == 'auto' else box.margin_right
@@ -116,16 +115,17 @@ def flex_layout(context, box, bottom_space, skip_stack, containing_block,
     # TODO: this will clobber any used width/height/margin attributes,
     # is that really what we want?
     resolve_percentages(parent_box, containing_block)
-    # TODO: removing auto margins is OK for this step, but margins should be
-    # calculated later.
+    # NOTE: we remove auto margins from parent_box (which is a
+    # throwaway) only but keep them on box itself.  They will get
+    # computed later, once we have done some layout.
     if parent_box.margin_top == 'auto':
-        box.margin_top = parent_box.margin_top = 0
+        parent_box.margin_top = 0
     if parent_box.margin_bottom == 'auto':
-        box.margin_bottom = parent_box.margin_bottom = 0
+        parent_box.margin_bottom = 0
     if parent_box.margin_left == 'auto':
-        box.margin_left = parent_box.margin_left = 0
+        parent_box.margin_left = 0
     if parent_box.margin_right == 'auto':
-        box.margin_right = parent_box.margin_right = 0
+        parent_box.margin_right = 0
     if isinstance(parent_box, boxes.FlexBox):
         block.block_level_width(parent_box, containing_block)
     else:
@@ -339,7 +339,7 @@ def flex_layout(context, box, bottom_space, skip_stack, containing_block,
     # using the rules of the formatting context in which it
     # participates. For this computation, auto margins on flex items
     # are treated as 0.) TODO: the whole step has to be fixed
-    if axis == 'width':
+    if axis == 'width':  # It needs to be a number from here on
         block.block_level_width(box, containing_block)
     else:
         if box.style['height'] != 'auto' and box.style['height'].unit != '%':
@@ -993,6 +993,10 @@ def flex_layout(context, box, bottom_space, skip_stack, containing_block,
 
     # Step 13 - 9.6. Cross-Axis Alignment
     LOGGER.debug("%r 9.6 Cross-Axis Alignment", box)
+    # Make sure width/margins are no longer "auto", as we did not do
+    # it above in step 9.2.4.
+    if cross == 'width':
+        block.block_level_width(box, containing_block)
     position_cross = (
         box.content_box_y() if cross == 'height'
         else box.content_box_x())