Fix typing errors in samples_per_second.py

Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Jun 20, 2024 · bc74464 · bc74464
1 parent 2c15cf5
commit bc74464
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/project/algorithms/callbacks/samples_per_second.py b/project/algorithms/callbacks/samples_per_second.py
@@ -1,7 +1,7 @@
 import time
 
 from lightning import LightningModule, Trainer
-from torch import Tensor, nn
+from torch import Tensor
 from torch.optim import Optimizer
 
 from project.algorithms.bases.algorithm import Algorithm, BatchType, StepOutputDict
@@ -19,7 +19,7 @@ def __init__(self):
     def on_shared_epoch_start(
         self,
         trainer: Trainer,
-        pl_module: Algorithm[BatchType, StepOutputDict, nn.Module],
+        pl_module: Algorithm[BatchType, StepOutputDict],
         phase: PhaseStr,
     ) -> None:
         self.last_update_time.clear()
@@ -34,7 +34,7 @@ def on_shared_epoch_start(
     def on_shared_batch_end(
         self,
         trainer: Trainer,
-        pl_module: Algorithm[BatchType, StepOutputDict, nn.Module],
+        pl_module: Algorithm[BatchType, StepOutputDict],
         outputs: StepOutputDict,
         batch: BatchType,
         batch_index: int,