Commit eea51ac6 authored by muyangli's avatar muyangli
Browse files

fix: fix the forward_layer functions for single-stream blocks

parent 67db2410
......@@ -924,7 +924,7 @@ std::tuple<Tensor, Tensor> FluxModel::forward_layer(size_t layer,
hidden_states, encoder_hidden_states, temb, rotary_emb_img, rotary_emb_context, 0.0f);
} else {
std::tie(hidden_states, encoder_hidden_states) =
transformer_blocks.at(layer - transformer_blocks.size())
single_transformer_blocks.at(layer - transformer_blocks.size())
->forward(hidden_states, encoder_hidden_states, temb, rotary_emb_img, rotary_emb_context, 0.0f);
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment