[BUGFIX] Pruning example works on GPU now

2021-06-03 14:35:24 +02:00
parent 1b09b1d57b
commit e209bf73d5
2 changed files with 9 additions and 11 deletions
--- a/prototorch/models/callbacks.py
+++ b/prototorch/models/callbacks.py
@@ -17,9 +17,6 @@ class PruneLoserPrototypes(pl.Callback):
        self.frequency = frequency
        self.verbose = verbose

-    def on_epoch_start(self, trainer, pl_module):
-        pl_module.initialize_prototype_win_ratios()
-
    def on_epoch_end(self, trainer, pl_module):
        if (trainer.current_epoch + 1) < self.idle_epochs:
            return None