React Grid Layout is not responsive when used inside a parent conatiner whose width is changing when screen size change

[viramjain13]

I am using a responsive grid layout inside the div which has some width which is changing when we change the screen size so when I drag the screen and pop back to the original size the grid column is not changing from 2 to 3 on the large screen according to the breakpoints
here is my code,

import { useMediaQuery, useTheme } from "@mui/material";
import React, { useEffect, useState } from "react";
import { Layout, Responsive, WidthProvider } from "react-grid-layout";

const ResponsiveGridLayout = WidthProvider(Responsive);

function App() {
const layout: Layout[] = [
{
i: "1",
x: 0,
y: 0,
h: 1,
w: 1,
},
{
i: "2",
x: 1,
y: 0,
h: 1,
w: 1,
},
{
i: "3",
x: 2,
y: 0,
h: 1,
w: 1,
},
{
i: "4",
x: 0,
y: 1,
h: 1,
w: 1,
},
];

const [test, setTest] = useState(layout);
const theme = useTheme();

let width = 330;

if (useMediaQuery(theme.breakpoints.up(400))) {
width = 372;
}

if (useMediaQuery(theme.breakpoints.up(950))) {
width = 760;
}

if (useMediaQuery(theme.breakpoints.up(1120))) {
width = 1148;
}

const [widthState, setWidthSate] = useState(width);

useEffect(() => {
setWidthSate(() => width);
}, [width]);

return (
<div style={{ width: widthState }}>
<ResponsiveGridLayout
style={{ backgroundColor: "red" }}
breakpoints={{ lg: 1340, md: 950, sm: 0 }}
cols={{ lg: 3, md: 2, sm: 1 }}
margin={{ lg: [16, 16], md: [16, 16], sm: [16, 16] }}
rowHeight={300}
layouts={{ lg: layout, md: layout }}
width={widthState}
useCSSTransforms
onLayoutChange={(layout, layouts) => setTest(layout)}
>
{test.map((m) => (
<div style={{ backgroundColor: "white" }} key={m.i}>
<div
style={{
backgroundColor: "blue",
}}
>

))}

);
}

export default App;
can someone please help me with that?
I really appreciate any help you can provide.