| | |
|---|
| | | ) |
|---|
| | | elif type == 'torchscripts': |
|---|
| | | device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|---|
| | | print("Exporting torchscripts on device {}".format(device)) |
|---|
| | | _torchscripts( |
|---|
| | | m, |
|---|
| | | path=export_dir, |
|---|
| | |
|---|
| | | 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}.torchscripts")) |
|---|
| | | model_script.save(os.path.join(path, f"{model.export_name}_blade.torchscripts")) |
|---|
