From dffc0f96ec76fc957a64be2ba828003e2ce09df9 Mon Sep 17 00:00:00 2001 From: Ivan Yashchuk Date: Thu, 30 May 2024 19:14:56 +0300 Subject: [PATCH] Update FlopCounterMode usage in throughput.py (#19926) `mods` argument is not needed anymore for `FlopCounterMode`: https://github.com/pytorch/pytorch/blob/ffe506e85350a505be5698c871d50b2fc614406d/torch/utils/flop_counter.py#L595-L596 --- src/lightning/fabric/utilities/throughput.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning/fabric/utilities/throughput.py b/src/lightning/fabric/utilities/throughput.py index c340686346..f483c274c3 100644 --- a/src/lightning/fabric/utilities/throughput.py +++ b/src/lightning/fabric/utilities/throughput.py @@ -296,7 +296,7 @@ def measure_flops( raise ImportError("`measure_flops` requires PyTorch >= 2.1.") from torch.utils.flop_counter import FlopCounterMode - flop_counter = FlopCounterMode(model, display=False) + flop_counter = FlopCounterMode(display=False) with flop_counter: if loss_fn is None: forward_fn()