Skip to content

Commit aded8bc

Browse files
github-actions[bot]mvafin
authored andcommitted
Apply style fixes
1 parent 7015abd commit aded8bc

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

optimum/exporters/openvino/__main__.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -360,17 +360,20 @@ class StoreAttr(object):
360360

361361
GPTQQuantizer.post_init_model = post_init_model
362362
if do_bitnet_patching:
363-
from transformers.integrations.bitnet import AutoBitLinear, unpack_weights
364363
import functools
365364

365+
from transformers.integrations.bitnet import AutoBitLinear, unpack_weights
366+
366367
orig_load_hook = AutoBitLinear.load_hook
367368

368369
# rewrite load hook to save original weight
369370
@functools.wraps(orig_load_hook)
370371
def bitnet_load_hook(self, state_dict, prefix, *args, **kwargs):
371372
if (prefix + "weight") in state_dict and state_dict[prefix + "weight"].dtype != self.weight.dtype:
372373
self.original_weight = state_dict[prefix + "weight"]
373-
state_dict[prefix + "weight"] = unpack_weights(state_dict[prefix + "weight"], dtype=self.weight.dtype).to(torch.device("meta"))
374+
state_dict[prefix + "weight"] = unpack_weights(
375+
state_dict[prefix + "weight"], dtype=self.weight.dtype
376+
).to(torch.device("meta"))
374377
return state_dict
375378

376379
AutoBitLinear.load_hook = bitnet_load_hook

0 commit comments

Comments
 (0)