perf: optimize OneHot shader, do not branch inside loop

wonnx/templates/endomorphism/onehot.wgsl

@@ -25,21 +25,22 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
 	let index_of_index = global_id.x;
 	let depth = u32(input_depth.data);
 	var index = input_indexes.data[index_of_index];
+
 	if(index < 0) {
 		index = i32(depth) + index;
 	}
 
 	let offset = index_of_index * depth;
 	let off_value = input_values.data[0];
-	let on_value = input_values.data[1];
 
 	for(var i = 0; i < i32(depth); i = i + 1) {
 		let ii = offset + u32(i);
-		if(i == index) {
-			output_0.data[ii] = on_value;
-		}
-		else {
-			output_0.data[ii] = off_value;
-		}
+		output_0.data[ii] = off_value;
+	}
+
+	if(index < i32(depth)) {
+		let on_value = input_values.data[1];
+		let hot_index = offset + u32(index);
+		output_0.data[hot_index] = on_value;
 	}
 }