typing

Borda · Borda · commit 41932ec59163 · 2025-04-22T22:09:29.000+02:00
diff --git a/src/lightning/pytorch/callbacks/throughput_monitor.py b/src/lightning/pytorch/callbacks/throughput_monitor.py
@@ -140,7 +140,7 @@ def _update(self, trainer: "Trainer", pl_module: "LightningModule", batch: Any,
             # this assumes that all iterations used the same batch size
             samples=iter_num * batch_size,
             lengths=None if self.length_fn is None else self._lengths[stage],
-            flops=flops_per_batch,
+            flops=int(flops_per_batch) if flops_per_batch is not None else None,
         )
 
     def _compute(self, trainer: "Trainer", iter_num: Optional[int] = None) -> None:

Original file line number	Diff line number	Diff line change
`@@ -140,7 +140,7 @@ def _update(self, trainer: "Trainer", pl_module: "LightningModule", batch: Any,`
`140`	`140`	`# this assumes that all iterations used the same batch size`
`141`	`141`	`samples=iter_num * batch_size,`
`142`	`142`	`lengths=None if self.length_fn is None else self._lengths[stage],`
`143`		`- flops=flops_per_batch,`
	`143`	`+ flops=int(flops_per_batch) if flops_per_batch is not None else None,`
`144`	`144`	`)`
`145`	`145`
`146`	`146`	`def _compute(self, trainer: "Trainer", iter_num: Optional[int] = None) -> None:`