diff --git a/drivers/base/class.c b/drivers/base/class.c
index 6e2c3b0..ccf312d 100644
--- a/drivers/base/class.c
+++ b/drivers/base/class.c
@@ -192,7 +192,6 @@ int __class_register(struct class *cls, struct lock_class_key *key)
 		return error;
 	}
 	error = add_class_attrs(class_get(cls));
-	class_put(cls);
 	return error;
 }
 EXPORT_SYMBOL_GPL(__class_register);
@@ -200,6 +199,7 @@ EXPORT_SYMBOL_GPL(__class_register);
 void class_unregister(struct class *cls)
 {
 	pr_debug("device class '%s': unregistering\n", cls->name);
+	class_put(cls);
 	remove_class_attrs(cls);
 	kset_unregister(&cls->p->class_subsys);
 }