[major] fix error when device_id != 0

88b96b2a · sxtyzhangzk · muyangli · 92d283a2 · 88b96b2a
Commit 88b96b2a authored Apr 17, 2025 by sxtyzhangzk Committed by muyangli Apr 19, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/FluxModel.cpp src/FluxModel.cpp +1 -1

No files found.
--- a/src/FluxModel.cpp
+++ b/src/FluxModel.cpp
@@ -774,7 +774,7 @@ FluxModel::FluxModel(bool use_fp4, bool offload, Tensor::ScalarType dtype, Devic
        }
    }
    for (int i = 0; i < 38; i++) {
-        single_transformer_blocks.push_back(std::make_unique<FluxSingleTransformerBlock>(3072, 24, 3072, 4, use_fp4, dtype, Device::cuda()));
+        single_transformer_blocks.push_back(std::make_unique<FluxSingleTransformerBlock>(3072, 24, 3072, 4, use_fp4, dtype, device));
        registerChildren(*single_transformer_blocks.back(), format("single_transformer_blocks.{}", i));
        if (offload) {
            single_transformer_blocks.back()->setLazyLoad(true);