Update throughput docs (#19415)
This commit is contained in:
parent
28a80238a4
commit
78b7a39e72
|
@ -211,6 +211,7 @@ Get Started
|
|||
Loggers <api/loggers>
|
||||
Precision <api/precision>
|
||||
Strategies <api/strategies>
|
||||
Utilities <api/utilities>
|
||||
|
||||
|
||||
.. toctree::
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue