fix dim

rwth-i6 · Aug 9, 2023 · 48f9e5b · 48f9e5b
1 parent 54bb95d
commit 48f9e5b
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/nn/naming.py b/nn/naming.py
@@ -1286,7 +1286,13 @@ def dim_repr(self):
             assert not dim.match_priority
             # We assume Dim is imported.
             # generic fallback
-            return f"Dim(kind={dim.kind}, description={dim.description!r}, dimension={dim.dimension})"
+            kind_s = {
+                None: "None",
+                nn.Dim.Types.Batch: "Dim.Types.Batch",
+                nn.Dim.Types.Spatial: "Dim.Types.Spatial",
+                nn.Dim.Types.Feature: "Dim.Types.Feature",
+            }[dim.kind]
+            return f"Dim(description={dim.description!r}, dimension={dim.dimension}, kind={kind_s})"
 
     class SetProxy:
         """