diff --git a/src/runtime_src/core/edge/drm/zocl/edge/zocl_aie.c b/src/runtime_src/core/edge/drm/zocl/edge/zocl_aie.c
index fe6b20c55e..c814d76e79 100644
--- a/src/runtime_src/core/edge/drm/zocl/edge/zocl_aie.c
+++ b/src/runtime_src/core/edge/drm/zocl/edge/zocl_aie.c
@@ -366,11 +366,8 @@ zocl_create_aie(struct drm_zocl_dev *zdev, struct axlf *axlf, char __user *xclbi
 	zdev->aie->uid = req.uid;
 
 	/* Register AIE error call back function. */
-	/* only aie-1 supports error management*/
-	if (hw_gen == 1) {
-		rval = aie_register_error_notification(zdev->aie->aie_dev,
-		zocl_aie_error_cb, zdev);
-	}
+	rval = aie_register_error_notification(zdev->aie->aie_dev,
+					       zocl_aie_error_cb, zdev);
 	mutex_unlock(&zdev->aie_lock);
 
 	zocl_init_aie(zdev);