Skip to content

Commit 66390f4

Browse files
guangy10malfet
authored andcommitted
Fix name_to_dtype error in export.py (#122)
1 parent 0166817 commit 66390f4

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

export.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323

2424
from model import Transformer
2525
from generate import _load_model, decode_one_token
26-
from quantize import quantize_model
26+
from quantize import quantize_model, name_to_dtype
2727
from torch._export import capture_pre_autograd_graph
2828

2929
default_device = "cpu" # 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -77,7 +77,7 @@ def main(checkpoint_path, device, quantize = "{ }", args = None):
7777
# dtype:
7878
if args.dtype:
7979
model.to(dtype=name_to_dtype(args.dtype))
80-
80+
8181
model = model_wrapper(model, device=device)
8282

8383
output_pte_path = args.output_pte_path

0 commit comments

Comments
 (0)