fix format

microsoft · May 30, 2024 · 416fae4 · 416fae4
1 parent ce6b694
commit 416fae4
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/deepspeed/inference/v2/modules/heuristics.py b/deepspeed/inference/v2/modules/heuristics.py
@@ -132,7 +132,7 @@ def instantiate_moe(moe_config: DSMoEConfig, engine_config: RaggedInferenceEngin
         }
 
     # check if we are on H100 or above
-    if torch.cuda.get_device_capability(0)[0] >= 9: #ignore-cuda
+    if torch.cuda.get_device_capability(0)[0] >= 9:  #ignore-cuda
         config = ConfigBundle(name="pytorch_multi_gemm_moe",
                               config=moe_config,
                               implementation_config=implementation_config)

diff --git a/deepspeed/inference/v2/modules/implementations/moe/pytorch_multi_gemm.py b/deepspeed/inference/v2/modules/implementations/moe/pytorch_multi_gemm.py
@@ -231,7 +231,6 @@ def forward(self,
             hidden_states (torch.Tensor): Hidden states tensor. Expected shape is [batch, seq_len, model_dim].
             gate_w (torch.Tensor): Gate weight tensor. Expected shape is [num_experts, model_dim].
         """
-        print("Using DSPytorchMultiGemmMoE forward pass")
 
         moe_input, expert_cumsum, scores, mapped_slots = self._gate(hidden_states, batch_metadata, gate_w)