blob: ae3fbcdca0995c27948f5bd288b554bee0f0f9fe (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
package gpt2
import ort "github.com/yalue/onnxruntime_go"
func SessionsOptionsWithCUDADeviceID(deviceID string) (*ort.SessionOptions, error) {
var sessionOptions *ort.SessionOptions
var cudaProviderOptions *ort.CUDAProviderOptions
if s, err := ort.NewSessionOptions(); err != nil {
return nil, err
} else {
sessionOptions = s
}
if c, err := ort.NewCUDAProviderOptions(); err != nil {
return nil, err
} else {
cudaProviderOptions = c
}
if err := cudaProviderOptions.Update(map[string]string{
"device_id": deviceID,
}); err != nil {
return nil, err
}
if err := sessionOptions.AppendExecutionProviderCUDA(cudaProviderOptions); err != nil {
return nil, err
}
if err := cudaProviderOptions.Destroy(); err != nil {
return nil, err
}
return sessionOptions, nil
}
|
