How can I introduce a tikz diagram in between two paragraphs?

Question

I am writing a proof which uses diagrams, I have some text that needs to go above and below a diagram. When I compile the following code, both pieces of text go above, and the diagram ends up below. What can I do so that the diagram ends up between both pieces of text?

Code is below:

\textbf{Proof:} Consider the following fractions of paths.
\begin{figure}
  \centering
  \begin{subfigure}{0.4\textwidth}
    \centering
    \begin{tikzpicture}
    \draw[step=1cm,black,very thick] (0,0) grid (3,3);
      % Draw the grid and color alternate squares
      \foreach \x in {0,1,2}
        \foreach \y in {0,1,2}
        {
          \pgfmathtruncatemacro\result{\x+\y}
          \ifodd\result
            \fill[black] (\x,\y) rectangle ++(1,1);
          \else
            \fill[white] (\x,\y) rectangle ++(1,1);
          \fi
        }
  % Draw the thicker and orange line path passing through the centers of all squares
  \draw[thick, orange] (0.5,1.5) -- (0.5,2.5);
\end{tikzpicture}
\caption{Edge to corner path}

\end{subfigure}
  \hspace{2cm} % Adjust the horizontal space between subfigures
  \begin{subfigure}{0.4\textwidth}
    \centering
    \begin{tikzpicture}
    \draw[step=1cm,black,very thick] (0,0) grid (3,3);
      % Draw the grid and color alternate squares
      \foreach \x in {0,1,2}
        \foreach \y in {0,1,2}
        {
          \pgfmathtruncatemacro\result{\x+\y}
          \ifodd\result
            \fill[black] (\x,\y) rectangle ++(1,1);
          \else
            \fill[white] (\x,\y) rectangle ++(1,1);
          \fi
        }
  % Draw the thicker and orange line path passing through the centers of all squares
  \draw[thick, orange] (0.5,1.5) -- (1.5,1.5);
\end{tikzpicture}
\caption{Edge to center path}

\end{subfigure}
  \caption{Edge cases}
\end{figure}
Subfigure (a) shows the path from the edge square to the center, and (b) shows the path from the edge square to an adjacent corner square.```