better str method to account for ablated neurons

mklasby · Aug 9, 2024 · a7bac88 · a7bac88
1 parent 2c09afb
commit a7bac88
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/sparsimony/dst/srigl.py b/sparsimony/dst/srigl.py
@@ -191,7 +191,10 @@ def __str__(self) -> str:
         for config in self.groups:
             mask = get_mask(**config)
             mask_flat = mask.view(mask.shape[0], prod(mask.shape[1:]))
-            ffi.append(mask_flat.sum(dim=1, dtype=torch.int).unique().item())
+            this_ffi = mask_flat.sum(dim=1, dtype=torch.int).unique()
+            if len(this_ffi) > 1:
+                this_ffi = this_ffi[this_ffi != 0]
+            ffi.append(this_ffi.item())
         s = super().__str__()
         s += f"FFI: {ffi}\n"
         return s