Fabrice-TIERCELIN
commited on
' instead of "
Browse files
hyvideo/modules/fp8_optimization.py
CHANGED
@@ -91,7 +91,7 @@ def convert_fp8_linear(module, dit_weight_path, original_dtype, params_to_keep={
|
|
91 |
|
92 |
fp8_layers = []
|
93 |
for key, layer in module.named_modules():
|
94 |
-
if isinstance(layer, nn.Linear) and (
|
95 |
fp8_layers.append(key)
|
96 |
original_forward = layer.forward
|
97 |
layer.weight = torch.nn.Parameter(layer.weight.to(torch.float8_e4m3fn))
|
|
|
91 |
|
92 |
fp8_layers = []
|
93 |
for key, layer in module.named_modules():
|
94 |
+
if isinstance(layer, nn.Linear) and ("double_blocks" in key or "single_blocks" in key):
|
95 |
fp8_layers.append(key)
|
96 |
original_forward = layer.forward
|
97 |
layer.weight = torch.nn.Parameter(layer.weight.to(torch.float8_e4m3fn))
|