cbensimon HF Staff commited on
Commit
495d59a
·
1 Parent(s): 9292d1e

Package AOTI

Browse files
Files changed (1) hide show
  1. app.py +3 -0
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: