qkv_out can be a single tensor or a list. Handling these cases separe…

…tely. (#1850) Co-authored-by: Jeff Rasley <[email protected]>
microsoft · Mar 24, 2022 · ebbcfd5 · ebbcfd5
1 parent c7af747
commit ebbcfd5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/ops/transformer/inference/transformer_inference.py b/deepspeed/ops/transformer/inference/transformer_inference.py
@@ -274,7 +274,7 @@ def selfAttention_fp():
                                    norm_b,
                                    config.epsilon,
                                    (attn_qkvb is not None))
-            context_layer, key_layer, value_layer = compute_attention(qkv_out[0], input_mask)
+            context_layer, key_layer, value_layer = compute_attention(qkv_out[0] if isinstance(qkv_out, list) else qkv_out, input_mask)
             output = vector_matmul_func(context_layer, attn_ow, False)
 
             return output, key_layer, value_layer, context_layer, qkv_out[-1] # attn_out, present_key, present_value, context_output, inp_norm