Spaces:
Running
on
Zero
Running
on
Zero
Package AOTI
Browse files
app.py
CHANGED
@@ -15,6 +15,7 @@ import gradio as gr
|
|
15 |
import spaces
|
16 |
import torch
|
17 |
import torch._inductor
|
|
|
18 |
from torch.export.pt2_archive._package import AOTICompiledModel
|
19 |
from torch.export.pt2_archive._package_weights import Weights
|
20 |
from torchvision.models import ResNet18_Weights, resnet18
|
@@ -41,6 +42,8 @@ def compile_model():
|
|
41 |
'aot_inductor.package': True,
|
42 |
'max_autotune': True,
|
43 |
})
|
|
|
|
|
44 |
weights, = (artifact for artifact in artifacts if isinstance(artifact, Weights))
|
45 |
weights_: dict[str, torch.Tensor] = {}
|
46 |
for name in weights:
|
|
|
15 |
import spaces
|
16 |
import torch
|
17 |
import torch._inductor
|
18 |
+
from torch._inductor.package import package_aoti
|
19 |
from torch.export.pt2_archive._package import AOTICompiledModel
|
20 |
from torch.export.pt2_archive._package_weights import Weights
|
21 |
from torchvision.models import ResNet18_Weights, resnet18
|
|
|
42 |
'aot_inductor.package': True,
|
43 |
'max_autotune': True,
|
44 |
})
|
45 |
+
files = [file for file in artifacts if isinstance(file, str)]
|
46 |
+
package_aoti(package_path, files)
|
47 |
weights, = (artifact for artifact in artifacts if isinstance(artifact, Weights))
|
48 |
weights_: dict[str, torch.Tensor] = {}
|
49 |
for name in weights:
|