chore: remove optimizer_idx from all steps

2023-10-25 15:03:13 +02:00
parent 60990f42d2
commit 71167a8f77
6 changed files with 14 additions and 19 deletions
--- a/src/prototorch/models/abstract.py
+++ b/src/prototorch/models/abstract.py
@@ -2,7 +2,6 @@

 import logging

-import prototorch
 import pytorch_lightning as pl
 import torch
 import torch.nn.functional as F
@@ -228,7 +227,7 @@ class NonGradientMixin(ProtoTorchMixin):
        super().__init__(*args, **kwargs)
        self.automatic_optimization = False

-    def training_step(self, train_batch, batch_idx, optimizer_idx=None):
+    def training_step(self, train_batch, batch_idx):
        raise NotImplementedError