feat: Add GPUType field for TensorRT (#107)

kleveross · Aug 21, 2020 · fd7f814 · fd7f814
1 parent 10ed03b
commit fd7f814
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/pkg/model/types.go b/pkg/model/types.go
@@ -17,6 +17,10 @@ type Metadata struct {
 	Tags            []string          `json:"tags,omitempty" yaml:"tags,omitempty"`
 	Labels          map[string]string `json:"labels,omitempty" yaml:"labels,omitempty"`
 	Format          string            `json:"format,omitempty" yaml:"format,omitempty"`
+	// GPUType is for TensorRT format only, it must be set when extract signature or serve 
+	// as a online service, otherwise， it can not extract or serve as a service.
+	// for other model format, you can set empty string or not set.
+	GPUType         string            `json:"gpuType,omitempty" yaml:"gpuType,omitempty"`
 	Framework       string            `json:"framework,omitempty" yaml:"framework,omitempty"`
 	Metrics         []Metric          `json:"metrics,omitempty" yaml:"metrics,omitempty"`
 	Hyperparameters []Hyperparameter  `json:"hyperparameters,omitempty" yaml:"hyperparameters,omitempty"`