Merge pull request #4 from seasonjs/patch-1

fix gpu error with darwin
seasonjs · Oct 23, 2023 · 706caa6 · 706caa6
2 parents b7746e1 + 2881930
commit 706caa6
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 2 deletions.
diff --git a/embed_darwin.go b/embed_darwin.go
@@ -19,6 +19,9 @@ var libRwkvArm []byte
 var libName = "librwkv-*.dylib"
 
 func getDl(gpu bool) []byte {
+	if gpu {
+		_, _ = GetGPUInfo()
+	}
 	if runtime.GOARCH == "amd64" {
 		return libRwkvAmd64
 	}

diff --git a/embed_windows.go b/embed_windows.go
@@ -32,11 +32,11 @@ var supportGpuTable = map[string][]byte{
 
 func getDl(gpu bool) []byte {
 	if gpu {
-		gpu, err := GetGPUInfo()
+		gpuInfo, err := GetGPUInfo()
 		if err != nil {
 			log.Println(err)
 		}
-		if supportGpuTable[gpu] != nil {
+		if supportGpuTable[gpuInfo] != nil {
 			return supportGpuTable[gpu]
 		} else {
 			log.Println("GPU not support, use CPU instead.")

diff --git a/gpu_darwin.go b/gpu_darwin.go
@@ -1,6 +1,8 @@
 // Copyright (c) seasonjs. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+//go:build darwin
+
 package rwkv
 
 func GetGPUInfo() (string, error) {

diff --git a/gpu_linux.go b/gpu_linux.go
@@ -1,6 +1,8 @@
 // Copyright (c) seasonjs. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+//go:build linux
+
 package rwkv
 
 func GetGPUInfo() (string, error) {