diff --git a/drivers/amlogic/i2c/aml_i2c.c b/drivers/amlogic/i2c/aml_i2c.c
index eb7d201ade61f0..384a97e311aaf3 100644
--- a/drivers/amlogic/i2c/aml_i2c.c
+++ b/drivers/amlogic/i2c/aml_i2c.c
@@ -1172,6 +1172,8 @@ static int aml_i2c_probe(struct platform_device *pdev)
     if(ret)
         printk(" class register i2c_class fail!\n");
 
+    platform_set_drvdata(pdev, i2c);
+
     return 0;
 }
 
@@ -1180,14 +1182,24 @@ static int aml_i2c_probe(struct platform_device *pdev)
 static int aml_i2c_remove(struct platform_device *pdev)
 {
     struct aml_i2c *i2c = platform_get_drvdata(pdev);
-    if (i2c->mode == I2C_INTERRUPT_MODE)
-	    free_irq(i2c->irq, i2c);
-    if (i2c->mode == I2C_TIMER_POLLING_MODE)
-    hrtimer_cancel(&i2c->aml_i2c_hrtimer);
-    mutex_destroy(i2c->lock);
-    i2c_del_adapter(&i2c->adap);
-    kzfree(i2c);
-    i2c= NULL;
+
+    if (i2c != NULL) {
+        if (i2c->mode == I2C_INTERRUPT_MODE) {
+            free_irq(i2c->irq, i2c);
+        }
+
+        if (i2c->mode == I2C_TIMER_POLLING_MODE) {
+            hrtimer_cancel(&i2c->aml_i2c_hrtimer);
+        }
+
+        mutex_destroy(i2c->lock);
+        i2c_del_adapter(&i2c->adap);
+        class_unregister(&i2c->cls);
+        kzfree(i2c);
+    } else {
+        printk(KERN_ALERT "[%s]: i2c is NULL\n", __func__);
+    }
+
     return 0;
 }