def METHOD_NAME(test_case, device):
    class MyModule(flow.nn.Module):
        def __init__(self):
            super().__init__()
            self.weight = flow.nn.Parameter(flow.tensor(5.0, device=device))
        def forward(self, x):
            return x * self.weight + 1.0
    my_module = MyModule()
    of_sgd = flow.optim.SGD(my_module.parameters(), lr=0.001, momentum=0.9)
    eager_out_list = []
    for i in range(3):
        x = flow.tensor(i * 1.0, device=device, requires_grad=False)
        of_eager_out = my_module(x)
        of_eager_out.backward()
        of_sgd.step()
        of_sgd.zero_grad()
        eager_out_list.append(of_eager_out)
    lazy_module = MyModule()
    class ScalarTrainGraph(flow.nn.Graph):
        def __init__(self):
            super().__init__()
            self.m = lazy_module
            of_sgd = flow.optim.SGD(lazy_module.parameters(), lr=0.001, momentum=0.9)
            # self.m = MyModule()
            # of_sgd = flow.optim.SGD(self.m.parameters(), lr=0.001, momentum=0.9)
            self.add_optimizer(of_sgd)
        def build(self, x):
            loss = self.m(x)
            loss.backward()
            return loss
    lazy_out_list = []
    scalar_g = ScalarTrainGraph()
    for i in range(3):
        x = flow.tensor(i * 1.0, device=device)
        of_lazy_out = scalar_g(x)
        lazy_out_list.append(of_lazy_out)
    for i in range(3):
        test_case.assertTrue(
            np.array_equal(lazy_out_list[i].numpy(), eager_out_list[i].numpy())
        )