diff --git a/src/ieu/bmu/byte.sv b/src/ieu/bmu/byte.sv
index efd43d84..2b98b83d 100644
--- a/src/ieu/bmu/byte.sv
+++ b/src/ieu/bmu/byte.sv
@@ -30,7 +30,8 @@
 `include "wally-config.vh"
 
 module byteUnit #(parameter WIDTH=32) (
-  input  logic [WIDTH-1:0] A, B,          // Operands
+  input  logic [WIDTH-1:0] A,             // Operands
+  input  logic ByteSelect,                // LSB of Immediate
   output logic [WIDTH-1:0] ByteResult);   // rev8, orcb result
 
   logic [WIDTH-1:0] OrcBResult, Rev8Result;
@@ -41,6 +42,6 @@ module byteUnit #(parameter WIDTH=32) (
     assign Rev8Result[WIDTH-i-1:WIDTH-i-8] = A[i+7:i];
   end
 
-  assign ByteResult = (B[0]) ? OrcBResult : Rev8Result;
+  assign ByteResult = (ByteSelect) ? OrcBResult : Rev8Result;
 
 endmodule
\ No newline at end of file
diff --git a/src/ieu/bmu/zbb.sv b/src/ieu/bmu/zbb.sv
index 3289d13a..13529d6a 100644
--- a/src/ieu/bmu/zbb.sv
+++ b/src/ieu/bmu/zbb.sv
@@ -54,7 +54,7 @@ module zbb #(parameter WIDTH=32) (
   logic [WIDTH-1:0] ExtResult;           // sign/zero extend result
 
   cnt #(WIDTH) cnt(.A(A), .B(B), .W64(W64), .CntResult(CntResult));
-  byteUnit #(WIDTH) bu(.A(A), .B(B), .ByteResult(ByteResult));
+  byteUnit #(WIDTH) bu(.A(A), .ByteSelect(B[0]), .ByteResult(ByteResult));
   ext #(WIDTH) ext(.A(A), .B(B), .ExtResult(ExtResult));
 
 
diff --git a/src/ieu/controller.sv b/src/ieu/controller.sv
index 7bbc8bf6..d27af3b4 100644
--- a/src/ieu/controller.sv
+++ b/src/ieu/controller.sv
@@ -244,7 +244,7 @@ module controller(
   assign ALUControlD = {W64D, SubArithD, ALUOpD};
 
   if (`ZBS_SUPPORTED | `ZBA_SUPPORTED | `ZBB_SUPPORTED | `ZBC_SUPPORTED) begin: bitmanipi //change the conditional expression to OR any Z supported flags
-    bmuctrl bmuctrl(.clk, .reset, .StallD, .FlushD, .InstrD, .ALUSelectD, .BSelectD, .ZBBSelectD, .BRegWriteD, BW64D, BALUOpD, .StallE, .FlushE, .ALUSelectE, .BSelectE, .ZBBSelectE);
+    bmuctrl bmuctrl(.clk, .reset, .StallD, .FlushD, .InstrD, .ALUSelectD, .BSelectD, .ZBBSelectD, .BRegWriteD, .BW64D, .BALUOpD, .StallE, .FlushE, .ALUSelectE, .BSelectE, .ZBBSelectE);
   end else begin: bitmanipi
     assign ALUSelectD = Funct3D;
     assign ALUSelectE = Funct3E;