feat: feature maps are now also considered, when checking for unprodu…

…ctive layers
MLRichter · Feb 10, 2022 · b43d7c4 · b43d7c4
1 parent 425a87f
commit b43d7c4
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/rfa_toolbox/vizualize.py b/rfa_toolbox/vizualize.py
@@ -24,12 +24,14 @@ def _feature_map_size_label(feature_map_size: Union[int, Sequence[int]]) -> str:
     if not isinstance(feature_map_size, Sequence) and not isinstance(
         feature_map_size, np.ndarray
     ):
-        return f"\\nFeature Map Res.: {feature_map_size} x {feature_map_size}"
-    else:
         return (
-            "\\nFeature Map Res.: "
-            f"{' x '.join(np.asarray(feature_map_size).astype(int).astype(str))}"
+            f"\\nFeature Map Res.: {min(feature_map_size, 1)} "
+            f"x {min(feature_map_size, 1)}"
         )
+    else:
+        fm = np.asarray(feature_map_size)
+        fm[fm < 1] = 1
+        return "\\nFeature Map Res.: " f"{' x '.join(fm.astype(int).astype(str))}"
 
 
 def visualize_node(