diff --git a/python/monarch/_src/actor/proc_mesh.py b/python/monarch/_src/actor/proc_mesh.py index ee340061..bf511419 100644 --- a/python/monarch/_src/actor/proc_mesh.py +++ b/python/monarch/_src/actor/proc_mesh.py @@ -139,7 +139,7 @@ def _labels(self) -> List[str]: def _new_with_shape(self, shape: Shape) -> "ProcMesh": device_mesh = ( None - if self._device_mesh is None + if self._maybe_device_mesh is None else self._device_mesh._new_with_shape(shape) ) return ProcMesh(self._proc_mesh, _mock_shape=shape, _device_mesh=device_mesh)