fix: gemm weight swapped dimension

the dimensions of the weights are swapped in some cases
ENOT-AutoDL · Mar 14, 2023 · 1626c74 · 1626c74
1 parent 73046a6
commit 1626c74
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnx2torch/node_converters/gemm.py b/onnx2torch/node_converters/gemm.py
@@ -70,7 +70,7 @@ def _(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
             if not trans_b:
                 weights = weights.T
 
-            in_features, out_features = weights.shape[0], weights.shape[1]
+            in_features, out_features = weights.shape[1], weights.shape[0]
             torch_module = nn.Linear(
                 in_features=in_features,
                 out_features=out_features,