diff --git a/npc/core/src/main/scala/Keyboard.scala b/npc/core/src/main/scala/Keyboard.scala
index 60c65c5..30159e0 100644
--- a/npc/core/src/main/scala/Keyboard.scala
+++ b/npc/core/src/main/scala/Keyboard.scala
@@ -91,7 +91,7 @@ class KeyboardSegController extends Module {
   val keycode_digits = VecInit(keycode(3,0)) ++ VecInit(keycode(7,4))
   val ascii = MuxLookup(keycode, 0.U)(keycode_to_ascii)
 
-  val seg_contoller = SegControllerGenerator(8)
+  val seg_contoller = Module(new SegControllerGenerator(8, UInt(8.W)))
 
   seg_contoller.io.in_segs := VecInit(Seq(keycode, ascii, counter.value, 0.U))
   io.segs := seg_contoller.io.segs
diff --git a/npc/core/src/main/scala/SegControllerGenerator.scala b/npc/core/src/main/scala/SegControllerGenerator.scala
index 49f05d3..7f2cb68 100644
--- a/npc/core/src/main/scala/SegControllerGenerator.scala
+++ b/npc/core/src/main/scala/SegControllerGenerator.scala
@@ -4,9 +4,9 @@ import chisel3._
 import chisel3.util._
 import chisel3.util.log2Ceil
 
-class SegControllerGenerator(seg_count: Int) extends Module {
+class SegControllerGenerator[T <: Data](seg_count: Int, t: T) extends Module {
   val io = IO(new Bundle {
-    val in_segs = Input(Vec(seg_count, UInt()))
+    val in_segs = Input(Vec(seg_count / ((t.getWidth + 3) / 4), t))
     val segs = Output(Vec(seg_count, UInt(8.W)))
   })
   val digit_to_seg = ((0 until 16).map(_.U)).zip(Seq(
@@ -15,17 +15,10 @@ class SegControllerGenerator(seg_count: Int) extends Module {
     "b00000001".U, "b00001001".U, "b00010001".U, "b11000001".U,
     "b01100011".U, "b10000101".U, "b01100001".U, "b01110001".U,
   ))
-  val vec_size = (io.in_segs.getWidth + 3) / 4
-  val vec = io.in_segs.asTypeOf(Vec(vec_size, UInt(4.W)))
+  val vec = io.in_segs.asTypeOf(Vec(seg_count, UInt(4.W)))
 
-  val seg_regs = RegInit(VecInit(Seq.fill(seg_count)(0.U(8.W))))
-  seg_regs := vec.map(MuxLookup(_, 0xFF.U)(digit_to_seg)) ++ Seq(0xFF.U, 0xFF.U)
+  val segs = VecInit(Seq.fill(seg_count)(0.U(8.W)))
+  segs := vec.map(MuxLookup(_, 0xFF.U)(digit_to_seg))
 
-  io.segs := seg_regs
-}
-
-object SegControllerGenerator {
-  def apply(seg_count: Int): SegControllerGenerator = {
-    new SegControllerGenerator(seg_count)
-  }
+  io.segs := segs
 }