.

quant.py

@@ -178,7 +178,7 @@ def FluxFp8GeMMProcessor(model: torch.nn.Module):
     )
     named_modules = list(model.named_modules())
     for name, linear in tqdm.tqdm(named_modules, desc="Quantizing weights to fp8"):
-        if isinstance(linear, torch.nn.Linear):
+        if isinstance(linear, torch.nn.Linear) and "blocks" in name:
             quant_weight, weight_scale = per_tensor_quantize(linear.weight)
             bias = copy.deepcopy(linear.bias) if linear.bias is not None else None
             quant_linear = FP8DynamicLinear(