bugfix in softmax exp equation

KULeuven-MICAS · Sep 5, 2024 · 3d3dc77 · 3d3dc77
1 parent 37b6cce
commit 3d3dc77
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/stream/classes/io/onnx/softmax.py b/stream/classes/io/onnx/softmax.py
@@ -120,10 +120,10 @@ def get_layer_node_user_format(self, input_shape: list[int], output_shape: list[
                 data["loop_dims"] = ["K", "C"]
             case 3:
                 data["equation"] = "O[b][k][c]+=I[b][k][c]+W[b][k]"
-                data["loop_dims"] = ["B", "C", "k"]
+                data["loop_dims"] = ["B", "K", "C"]
             case 4:
                 data["equation"] = "O[b][h][k][c]+=I[b][h][k][c]+W[b][h][k]"
-                data["loop_dims"] = ["B", "H", "C", "k"]
+                data["loop_dims"] = ["B", "H", "K", "C"]
             case _:
                 raise NotImplementedError