From 91040ec2f2d5f826dfee979fe11a5015a3a12c32 Mon Sep 17 00:00:00 2001
From: henryruhs <info@henryruhs.com>
Date: Thu, 30 Nov 2023 23:15:08 +0100
Subject: [PATCH] Fix install script

---
 install.js | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/install.js b/install.js
index ee7b477..2117123 100644
--- a/install.js
+++ b/install.js
@@ -10,13 +10,13 @@ function install(kernel)
 		}
 		return 'python install.py --onnxruntime default --torch default';
 	}
-	if (gpu === 'nvidia')
+	if ([ 'linux', 'win32' ].includes(platform) && gpu === 'nvidia')
 	{
 		return 'python install.py --onnxruntime cuda --torch cuda';
 	}
-	if (gpu === 'amd')
+	if (platform === 'linux' && gpu === 'amd')
 	{
-		return 'python install.py --onnxruntime default --torch cpu';
+		return 'python install.py --onnxruntime directml --torch rocm';
 	}
 	return 'python install.py --onnxruntime default --torch cpu';
 }
@@ -31,14 +31,9 @@ module.exports = async kernel =>
 				type: 'conda',
 				name:
 				[
-					'cmake',
 					'ffmpeg'
 				],
 				args: '-c conda-forge'
-			},
-			{
-				gpu: 'nvidia',
-				name: 'cuda'
 			}
 		],
 		run: