diff --git a/drivers/gpu/drm/imx/dpu/dpu-kms.c b/drivers/gpu/drm/imx/dpu/dpu-kms.c
index 50b11bdd176faa51a4cb24bba71222b4d1d06f54..4a9dd1396a7fa326e5b9710a81a0ac65d1d075ea 100644
--- a/drivers/gpu/drm/imx/dpu/dpu-kms.c
+++ b/drivers/gpu/drm/imx/dpu/dpu-kms.c
@@ -761,9 +761,16 @@ static int dpu_drm_atomic_check(struct drm_device *dev,
 		}
 
 		kfree(states);
+	};
+
+	drm_for_each_crtc(crtc, dev) {
+		if (pipe_states_prone_to_put[drm_crtc_index(crtc)]) {
+			crtc_state = drm_atomic_get_crtc_state(state, crtc);
+			if (WARN_ON(IS_ERR(crtc_state)))
+				return PTR_ERR(crtc_state);
 
-		if (pipe_states_prone_to_put[drm_crtc_index(crtc)])
 			dpu_atomic_put_possible_states_per_crtc(crtc_state);
+		}
 	}
 
 	ret = drm_atomic_helper_check_planes(dev, state);