Class Foo

fx/profiling_tracer.py:16–19

# Setup: a module with `record_function`
class Foo(torch.nn.Module):
  def forward(self, x):
    with torch.profiler.record_function('foo'):
      return torch.relu(x)

f = Foo()
x = torch.randn(5, 3, 2)