1 files changed, 36 insertions, 0 deletions
diff --git a/gpt2/cuda.go b/gpt2/cuda.go
new file mode 100644
index 0000000..ae3fbcd
--- /dev/null
+++ b/gpt2/cuda.go
@@ -0,0 +1,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
+}