Update throughput docs (#19415)

This commit is contained in:
Carlos Mocholí 2024-02-06 22:26:10 +01:00 committed by GitHub
parent 28a80238a4
commit 78b7a39e72
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 1 deletions

View File

@ -211,6 +211,7 @@ Get Started
Loggers <api/loggers>
Precision <api/precision>
Strategies <api/strategies>
Utilities <api/utilities>
.. toctree::

View File

@ -50,7 +50,7 @@ class ThroughputMonitor(Callback):
model = MyModel()
def sample_forward():
batch = torch.randn(...)
batch = torch.randn(..., device="meta")
return model(batch)
self.flops_per_batch = measure_flops(model, sample_forward, loss_fn=torch.Tensor.sum)