8 lines
221 B
Python
8 lines
221 B
Python
import torch
|
|
from src.export_to_onnx import export_to_onnx
|
|
from src.config.presets import SMALL
|
|
|
|
|
|
if __name__ == "__main__":
|
|
device = torch.device("cuda")
|
|
export_to_onnx(SMALL, "src/pretrained/amt_s.onnx", device) |