def METHOD_NAME(self):
    class M(torch.nn.Module):
        def __init__(self):
            super().__init__()
        def forward(self, x, weight, bias):
            x = torch.nn.functional.linear(x, weight, bias)
            x = torch.nn.functional.linear(x, weight, bias)
            x = torch.nn.functional.relu(x)
            x = torch.nn.functional.linear(x, weight, bias)
            x = torch.nn.functional.linear(x, weight, bias)
            x = torch.nn.functional.relu(x)
            return x
    inputs = (torch.randn(1, 5), torch.rand((5, 5)), torch.zeros(5))
    gm, _ = torch._dynamo.export(M(), aten_graph=True)(*inputs)
    gm.graph.eliminate_dead_code()
    module_partitions = get_source_partitions(gm.graph, [torch.nn.functional.linear, torch.nn.functional.relu])
    self.assertEqual(len(module_partitions), 2)
    self.assertEqual(len(module_partitions[torch.nn.functional.linear]), 4)
    self.assertEqual(len(module_partitions[torch.nn.functional.relu]), 2)