opt1 = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    opt2 = ZeroRedundancyOptimizer(
        model.parameters(),
        torch.optim.SGD,
        lr=0.01,
        momentum=0.9,
        grad_clipping=False)

    opt1.step()
    opt2.step()
    s1 = opt1.state_dict()
    s2 = opt2.state_dict()
    print("AFTER STEPPING ONCE")
    print("opt1.state", opt1.state)
    print("opt1.state_dict()", s1)
    print("opt2.state[base]", opt2.state['base'])
    print("opt2.state_dict()[base]", s2['base'])
    self.assertEqual(s1, s2['base'])

    # s1_clone = deepcopy(s1)
    # s2_clone = deepcopy(s2)
    opt1.load_state_dict(s1)
    opt2.load_state_dict(s2)
    print("AFTER LOADING THE STATE_DICTs, should be same as before")
    print("opt1.state", opt1.state)
    print("opt1.state_dict()", opt1.state_dict())
    print("opt2.state", opt2.state['base'])
    print("opt2.state_dict()[base]", opt2.state_dict()['base'])
    self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])

    # step still runnable
    opt1.step()
    opt2.step()
    print("AFTER STEPPING AGAIN, WILL be different")
    print("opt1.state", opt1.state)
    print("opt1.state_dict()", opt1.state_dict())
    print("opt2.state", opt2.state['base'])
    print("opt2.state_dict()[base]", opt2.state_dict()['base'])
    opt1.load_state_dict(s1)
    opt2.load_state_dict(s2)
    print("AFTER LOADING THE STATE_DICTs, should be same as before")
    print("opt1.state", opt1.state)
    print("opt1.state_dict()", opt1.state_dict())
    print("opt2.state", opt2.state['base'])
    print("opt2.state_dict()[base]", opt2.state_dict()['base'])
    self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])