Update zero_bubble_pp.py

hpcaitech · Sep 12, 2024 · 9802a7d · 9802a7d
1 parent 24727de
commit 9802a7d
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/colossalai/pipeline/schedule/zero_bubble_pp.py b/colossalai/pipeline/schedule/zero_bubble_pp.py
@@ -484,11 +484,7 @@ def backward_b_step(
             # loss backward; output_obj is loss; so output_obj_grad should be None
             assert output_obj_grad is None
 
-        if "hidden_states" in input_obj.keys():
-            input_obj_ = input_obj["hidden_states"]
-        else:
-            input_obj_ = input_obj["input_ids"]
-
+        input_obj_ = input_obj["hidden_states"]
         if output_obj_grad is None:
             optimizer.backward(output_obj, inputs=input_obj_, retain_graph=True)
         else: