It starts working if vec.shape is [768, 1] instead of [768]

AXKuhta · Aug 23, 2022 · 75ad160 · 75ad160
1 parent d000d73
commit 75ad160
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/RWKV-v4/index.html b/RWKV-v4/index.html
@@ -30,7 +30,7 @@
 				vec_d.fill(5)
 
 				const mat = new ort.Tensor('float32', mat_d, [768, 768])
-				const vec = new ort.Tensor('float32', vec_d, [768])
+				const vec = new ort.Tensor('float32', vec_d, [768, 1])
 
 				// prepare feeds. use model input names as keys.
 				var feeds = { mat: mat, vec: vec }

diff --git a/RWKV-v4/matmul_test.py b/RWKV-v4/matmul_test.py
@@ -11,7 +11,7 @@ def export():
 	model = MATMUL()
 
 	mat = torch.rand([768, 768])
-	vec = torch.rand([768])
+	vec = torch.rand([768, 1])
 
 	torch.onnx.export(model, args=(mat, vec), f="matmul.onnx", input_names = ["mat", "vec"], output_names = ["x"], verbose=True)