[BUGFIX] Fix RSLVQ

2021-06-01 17:44:10 +02:00
parent 9c1a41997b
commit 21023a88d7
2 changed files with 14 additions and 12 deletions
--- a/prototorch/models/abstract.py
+++ b/prototorch/models/abstract.py
@@ -4,6 +4,10 @@ from torch.optim.lr_scheduler import ExponentialLR


 class AbstractPrototypeModel(pl.LightningModule):
+    @property
+    def num_prototypes(self):
+        return len(self.proto_layer.components)
+
    @property
    def prototypes(self):
        return self.proto_layer.components.detach().cpu()