diff --git a/keyboards/yanghu/unicorne/unicorne.c b/keyboards/yanghu/unicorne/unicorne.c
index 27b87253d..bce8f5207 100644
--- a/keyboards/yanghu/unicorne/unicorne.c
+++ b/keyboards/yanghu/unicorne/unicorne.c
@@ -14,7 +14,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 #include "unicorne.h"
-
+#include "i2c_master.h"
 
 // Custom i2c init to enable internal pull up resistor for i2c.
 void i2c_init(void) {
@@ -23,13 +23,13 @@ void i2c_init(void) {
         is_initialised = true;
 
         // Try releasing special pins for a short time
-        palSetPadMode(I2C1_SCL_BANK, I2C1_SCL, PAL_MODE_INPUT);
-        palSetPadMode(I2C1_SDA_BANK, I2C1_SDA, PAL_MODE_INPUT);
+        palSetLineMode(I2C1_SCL_PIN, PAL_MODE_INPUT);
+        palSetLineMode(I2C1_SDA_PIN, PAL_MODE_INPUT);
 
         chThdSleepMilliseconds(10);
         // Use internal pull up since we do not have pull up on i2c pins in v1 design.
-        palSetPadMode(I2C1_SCL_BANK, I2C1_SCL, PAL_MODE_ALTERNATE(I2C1_SCL_PAL_MODE) | PAL_STM32_OTYPE_OPENDRAIN | PAL_STM32_PUPDR_PULLUP);
-        palSetPadMode(I2C1_SDA_BANK, I2C1_SDA, PAL_MODE_ALTERNATE(I2C1_SDA_PAL_MODE) | PAL_STM32_OTYPE_OPENDRAIN | PAL_STM32_PUPDR_PULLUP);
+        palSetLineMode(I2C1_SCL_PIN, PAL_MODE_ALTERNATE(I2C1_SCL_PAL_MODE) | PAL_OUTPUT_TYPE_OPENDRAIN | PAL_STM32_PUPDR_PULLUP);
+        palSetLineMode(I2C1_SDA_PIN, PAL_MODE_ALTERNATE(I2C1_SDA_PAL_MODE) | PAL_OUTPUT_TYPE_OPENDRAIN | PAL_STM32_PUPDR_PULLUP);
     }
 }
 
@@ -58,7 +58,7 @@ led_config_t g_led_config = {{
   {208, 44},
   {130, 60}
 }, {// LED Index to Flag
-  LED_FLAG_ALL, LED_FLAG_ALL, LED_FLAG_ALL, LED_FLAG_ALL, LED_FLAG_ALL, 
+  LED_FLAG_ALL, LED_FLAG_ALL, LED_FLAG_ALL, LED_FLAG_ALL, LED_FLAG_ALL,
   LED_FLAG_ALL, LED_FLAG_ALL, LED_FLAG_ALL
 }};
 // clang-format on
@@ -138,11 +138,13 @@ void set_keylog(uint16_t keycode, keyrecord_t *record) {
 __attribute__((weak)) void oled_render_keylog(void) { oled_write(keylog_str, false); }
 
 // Keymaps can override this function
-__attribute__((weak)) void oled_task_user(void) {
+__attribute__((weak)) bool oled_task_kb(void) {
+    if (!oled_task_user()) { return false; }
     /* oled_render_keylog(); */
     oled_render_layer();
     oled_render_mods();
     led_t led_state = host_keyboard_led_state();
     oled_render_capslock(led_state.caps_lock);
+    return true;
 }
 #endif