Use platform specific defaults for CUDA device ID

1 files changed, 13 insertions, 1 deletions

diff --git a/gpt2/model.go b/gpt2/model.go
index 5fef956..feea423 100644
--- a/gpt2/model.go
+++ b/gpt2/model.go
@@ -20,7 +20,9 @@ type Model struct {
 	allocator *Allocator
 }
 
-func NewModel(name string, deviceID string, cfg Config, opts Options) *Model {
+func NewModel(name string, cfg Config, opts Options) *Model {
+	deviceID := CUDADeviceID()
+
 	return &Model{
 		name:     name,
 		deviceID: deviceID,
@@ -29,6 +31,16 @@ func NewModel(name string, deviceID string, cfg Config, opts Options) *Model {
 	}
 }
 
+func CUDADeviceID() string {
+	id, ok := os.LookupEnv("ONNXRUNTIME_CUDA_DEVICE_ID")
+
+	if !ok {
+		return defaultDeviceID
+	}
+
+	return id
+}
+
 func SharedLibraryPath() string {
 	p, ok := os.LookupEnv("ONNXRUNTIME_SHARED_LIBRARY_PATH")