| | |
|---|
| | | try: |
|---|
| | | import torch_blade |
|---|
| | | except Exception as e: |
|---|
| | | print(f"failed to load torch_blade: {e}") |
|---|
| | | print(f"Warning, if you are exporting bladedisc, please install it and try it again: pip install -U torch_blade\n") |
|---|
| | | |
|---|
| | | |
|---|
| | | def export(model, data_in=None, quantize: bool = False, opset_version: int = 14, type='onnx', **kwargs): |
|---|
| | |
|---|
| | | model.encoder = _bladedisc_opt(model.encoder, input_data[:2]) |
|---|
| | | model.decoder = _bladedisc_opt(model.decoder, tuple(decoder_inputs)) |
|---|
| | | model_script = torch.jit.trace(model, input_data) |
|---|
| | | model_script.save(os.path.join(path, f"{model.export_name}_blade.torchscripts")) |
|---|
| | | model_script.save(os.path.join(path, f"{model.export_name}_blade.torchscripts")) |
|---|
