def METHOD_NAME(cls, path: str) -> None:
    rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    if torch.cuda.is_available():
        device: torch.device = torch.device(f"cuda:{rank}")
        backend = "nccl"
        torch.cuda.set_device(device)
    else:
        device: torch.device = torch.device("cpu")
        backend = "gloo"
    dist.init_process_group(backend=backend)
    num_float_features = 32
    tables = [
        EmbeddingBagConfig(
            num_embeddings=(i + 1) * 10,
            embedding_dim=(i + 1) * 4 * world_size,
            name="table_" + str(i),
            feature_names=["feature_" + str(i)],
        )
        for i in range(3)
    ]
    weighted_tables = [
        EmbeddingBagConfig(
            num_embeddings=(i + 1) * 10,
            embedding_dim=(i + 1) * 4 * world_size,
            name="weighted_table_" + str(i),
            feature_names=["weighted_feature_" + str(i)],
        )
        for i in range(2)
    ]
    m = TestSparseNN(
        tables=tables,
        num_float_features=num_float_features,
        weighted_tables=weighted_tables,
        dense_device=device,
    )
    m.sparse.ebc = trec_shard(
        module=m.sparse.ebc,
        device=device,
        plan=column_wise(ranks=list(range(world_size))),
    )
    m.sparse.weighted_ebc = trec_shard(
        module=m.sparse.weighted_ebc,
        device=device,
        plan=column_wise(ranks=list(range(world_size))),
    )
    m.over = replicate(m.over)
    m.dense = replicate(m.dense)
    ######## run one iteration ########
    _, local_batch = ModelInput.generate(
        batch_size=8,
        world_size=world_size,
        num_float_features=num_float_features,
        tables=tables,
        weighted_tables=weighted_tables,
    )
    batch = local_batch[0].to(device)
    m(batch)[1].sum().backward()
    state_dict = m.state_dict()
    writer = FileSystemWriter(path=path)
    reader = FileSystemReader(path=path)
    save_state_dict(state_dict, writer)
    p_sum = torch.zeros(1, device=device)
    for p in m.parameters():
        with torch.no_grad():
            if isinstance(p, ShardedTensor):
                if not p.local_shards():
                    continue
                p = p.local_tensor()
            p_sum += p.sum()
            p.zero_()
            assert p.sum() == 0
    load_state_dict(state_dict, reader)
    m.load_state_dict(state_dict)
    p_sum_loaded = torch.zeros(1, device=device)
    for p in m.parameters():
        with torch.no_grad():
            if isinstance(p, ShardedTensor):
                if not p.local_shards():
                    continue
                p = p.local_tensor()
            p_sum_loaded += p.sum()
    # TODO: debug why failing on OSS
    # assert p_sum.allclose(p_sum_loaded)