gpt2/cuda.go


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

package gpt2

import ort "github.com/yalue/onnxruntime_go"

func WithCUDAProvider(options *ort.SessionOptions, deviceID string) error {
	var cudaProviderOptions *ort.CUDAProviderOptions

	if c, err := ort.NewCUDAProviderOptions(); err != nil {
		return err
	} else {
		cudaProviderOptions = c
	}

	if err := cudaProviderOptions.Update(map[string]string{
		"device_id": deviceID,
	}); err != nil {
		return err
	}

	if err := options.AppendExecutionProviderCUDA(cudaProviderOptions); err != nil {
		return err
	}

	if err := cudaProviderOptions.Destroy(); err != nil {
		return err
	}

	return nil
}