diff --git a/kmk/extensions/oled.py b/kmk/extensions/oled.py
index 9d83a93..fcb7de3 100644
--- a/kmk/extensions/oled.py
+++ b/kmk/extensions/oled.py
@@ -147,8 +147,8 @@ class Oled(Extension):
             self.updateOLED(sandbox)
         return
 
-    def after_matrix_scan(self, keyboard):
-        if keyboard.matrix_update or keyboard.secondary_matrix_update:
+    def after_matrix_scan(self, sandbox):
+        if sandbox.matrix_update or sandbox.secondary_matrix_update:
             self.timer_time_reset()
             return
 
@@ -197,4 +197,4 @@ class Oled(Extension):
             if self._display.brightness > 0.5:
                 self._display.brightness = 0.5
         else: self._display.brightness = (self._brightness)
-        return
\ No newline at end of file
+        return