Function equal_state_dict

tests/test_utils.py:886–893

assert equal state_dict (for the shared keys between st_1 and st_2).

(st_1, st_2)

def equal_state_dict(st_1, st_2):
    """
    assert equal state_dict (for the shared keys between st_1 and st_2).
    """
    for key_st_1, val_st_1 in st_1.items():
        if key_st_1 in st_2:
            val_st_2 = st_2.get(key_st_1)
            assert_allclose(val_st_1, val_st_2)


TEST_TORCH_TENSORS: tuple = (torch.as_tensor,)