From d35f5fad1d26fc09cfd98f6e7faf17d6a48e4e8f Mon Sep 17 00:00:00 2001
From: Jonas Knobloch <jonas.knobloch@t-online.de>
Date: Thu, 23 Apr 2026 20:14:26 +0200
Subject: Use platform specific defaults for CUDA device ID

---
 gpt2/model.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

(limited to 'gpt2/model.go')

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")
 
-- 
cgit v1.3.1