Update yolo_bricks.py

MaxSigmoidAttnBlock Enhancement
AILab-CVC · Jun 5, 2024 · f5b3623 · f5b3623
1 parent 4d90f45
commit f5b3623
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/yolo_world/models/layers/yolo_bricks.py b/yolo_world/models/layers/yolo_bricks.py
@@ -39,7 +39,7 @@ def __init__(self,
                 embed_channels % num_heads == 0), \
             'out_channels and embed_channels should be divisible by num_heads.'
         self.num_heads = num_heads
-        self.head_channels = out_channels // num_heads
+        self.head_channels = embed_channels // num_heads
         self.use_einsum = use_einsum
 
         self.embed_conv = ConvModule(