diff --git a/drivers/char/tpm/tpm_ibmvtpm.c b/drivers/char/tpm/tpm_ibmvtpm.c
index 946025a7413b6b..ced6b9f0008dc2 100644
--- a/drivers/char/tpm/tpm_ibmvtpm.c
+++ b/drivers/char/tpm/tpm_ibmvtpm.c
@@ -294,6 +294,8 @@ static int tpm_ibmvtpm_remove(struct vio_dev *vdev)
 		kfree(ibmvtpm->rtce_buf);
 	}
 
+	/* For tpm_ibmvtpm_get_desired_dma */
+	dev_set_drvdata(&vdev->dev, NULL);
 	kfree(ibmvtpm);
 
 	return 0;
@@ -309,15 +311,16 @@ static int tpm_ibmvtpm_remove(struct vio_dev *vdev)
 static unsigned long tpm_ibmvtpm_get_desired_dma(struct vio_dev *vdev)
 {
 	struct tpm_chip *chip = dev_get_drvdata(&vdev->dev);
-	struct ibmvtpm_dev *ibmvtpm = dev_get_drvdata(&chip->dev);
+	struct ibmvtpm_dev *ibmvtpm;
 
 	/* ibmvtpm initializes at probe time, so the data we are
 	* asking for may not be set yet. Estimate that 4K required
 	* for TCE-mapped buffer in addition to CRQ.
 	*/
-	if (!ibmvtpm)
+	if (!chip)
 		return CRQ_RES_BUF_SIZE + PAGE_SIZE;
 
+	ibmvtpm = dev_get_drvdata(&chip->dev);
 	return CRQ_RES_BUF_SIZE + ibmvtpm->rtce_size;
 }