black is black

vince62s · Dec 19, 2023 · ea900d8 · ea900d8
1 parent 3f81b8f
commit ea900d8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/onmt/modules/rmsnorm.py b/onmt/modules/rmsnorm.py
@@ -26,12 +26,12 @@ def __init__(self, hidden_size: int, eps: float = 1e-6):
     def forward(self, hidden_states):
         if AWQ_INFERENCE_ENGINE:
             output = torch.empty_like(hidden_states)
-            if hidden_states.dim() == 2: # patch for multi experts
+            if hidden_states.dim() == 2:  # patch for multi experts
                 hidden_states = hidden_states.unsqueeze(0)
             awq_inference_engine.layernorm_forward_cuda(
                 hidden_states, self.weight, output, self.eps
             )
-            if hidden_states.dim() == 2: # patch for multi experts
+            if hidden_states.dim() == 2:  # patch for multi experts
                 output = output.unsqueeze(0)
             return output
         else: