diff --git a/drivers/usb/typec/ucsi/displayport.c b/drivers/usb/typec/ucsi/displayport.c
index 048381c058a5..1a5c76907827 100644
--- a/drivers/usb/typec/ucsi/displayport.c
+++ b/drivers/usb/typec/ucsi/displayport.c
@@ -263,10 +263,12 @@ static void ucsi_displayport_work(struct work_struct *work)
 	mutex_unlock(&dp->con->lock);
 }
 
-void ucsi_displayport_remove_partner(struct typec_altmode *alt)
+void ucsi_displayport_remove_partner(struct typec_altmode *pdev)
 {
+	const struct typec_altmode *alt;
 	struct ucsi_dp *dp;
 
+	alt = typec_altmode_get_partner(pdev);
 	if (!alt)
 		return;
 
diff --git a/drivers/usb/typec/ucsi/ucsi.c b/drivers/usb/typec/ucsi/ucsi.c
index ddf2ad3752de..921499407403 100644
--- a/drivers/usb/typec/ucsi/ucsi.c
+++ b/drivers/usb/typec/ucsi/ucsi.c
@@ -432,8 +432,12 @@ static int ucsi_register_altmodes(struct ucsi_connector *con, u8 recipient)
 		command |= UCSI_GET_ALTMODE_CONNECTOR_NUMBER(con->num);
 		command |= UCSI_GET_ALTMODE_OFFSET(i);
 		len = ucsi_run_command(con->ucsi, command, alt, sizeof(alt));
-		if (len <= 0)
+		if (len <= 0) {
+			/* Assuming that all alt modes are now registered */
+			if (recipient == UCSI_RECIPIENT_SOP && i)
+				break;
 			return len;
+		}
 
 		/*
 		 * This code is requesting one alt mode at a time, but some PPMs
@@ -464,9 +468,8 @@ static int ucsi_register_altmodes(struct ucsi_connector *con, u8 recipient)
 
 static void ucsi_unregister_altmodes(struct ucsi_connector *con, u8 recipient)
 {
-	const struct typec_altmode *pdev;
 	struct typec_altmode **adev;
-	int i = 0;
+	int i;
 
 	switch (recipient) {
 	case UCSI_RECIPIENT_CON:
@@ -479,16 +482,15 @@ static void ucsi_unregister_altmodes(struct ucsi_connector *con, u8 recipient)
 		return;
 	}
 
-	while (adev[i]) {
+	for (i = 0; adev[i]; i++) {
 		if (recipient == UCSI_RECIPIENT_SOP &&
 		    (adev[i]->svid == USB_TYPEC_DP_SID ||
-			(adev[i]->svid == USB_TYPEC_NVIDIA_VLINK_SID &&
-			adev[i]->vdo != USB_TYPEC_NVIDIA_VLINK_DBG_VDO))) {
-			pdev = typec_altmode_get_partner(adev[i]);
-			ucsi_displayport_remove_partner((void *)pdev);
-		}
+		     (adev[i]->svid == USB_TYPEC_NVIDIA_VLINK_SID &&
+		      adev[i]->vdo != USB_TYPEC_NVIDIA_VLINK_DBG_VDO)))
+			ucsi_displayport_remove_partner(adev[i]);
+
 		typec_unregister_altmode(adev[i]);
-		adev[i++] = NULL;
+		adev[i] = NULL;
 	}
 }
 
@@ -583,8 +585,8 @@ static void ucsi_partner_change(struct ucsi_connector *con)
 	ret = ucsi_register_altmodes(con, UCSI_RECIPIENT_SOP);
 	if (ret)
 		dev_err(con->ucsi->dev,
-			"con%d: failed to register partner alternate modes\n",
-			con->num);
+			"con%d: failed to register partner alternate modes (%d)\n",
+			con->num, ret);
 	else
 		ucsi_altmode_update_active(con);
 }
diff --git a/include/linux/usb/typec_altmode.h b/include/linux/usb/typec_altmode.h
index d834e236c6df..85267df13bf1 100644
--- a/include/linux/usb/typec_altmode.h
+++ b/include/linux/usb/typec_altmode.h
@@ -40,7 +40,8 @@ static inline void typec_altmode_set_drvdata(struct typec_altmode *altmode,
 	dev_set_drvdata(&altmode->dev, data);
 }
 
-static inline void *typec_altmode_get_drvdata(struct typec_altmode *altmode)
+static inline void *
+typec_altmode_get_drvdata(const struct typec_altmode *altmode)
 {
 	return dev_get_drvdata(&altmode->dev);
 }