def METHOD_NAME(self):
    for preferred_nnlib in ["sklearn", "faiss"]:
        train_x = torch.arange(5, dtype=torch.float32).unsqueeze(-1)
        nn_util = NNUtil(k=3, dim=train_x.size(-1), preferred_nnlib=preferred_nnlib)
        # set train data
        nn_util.set_nn_idx(train_x)
        nn_idx = nn_util.find_nn_idx(train_x)
        nn_idx = torch.sort(nn_idx, dim=-1)[0]
        expected_idx = torch.tensor([[0, 1, 2], [0, 1, 2], [1, 2, 3], [2, 3, 4], [2, 3, 4]], dtype=torch.int64)
        assert torch.equal(
            nn_idx, expected_idx
        ), f"Preferred_nnlib = {preferred_nnlib}, nn_idx={nn_idx}, expected_idx={expected_idx}"
        # single point test
        test_x = torch.tensor([20], dtype=torch.float32)  # shape would be automatically expanded
        nn_idx = nn_util.find_nn_idx(test_x)
        nn_idx = torch.sort(nn_idx, dim=-1)[0]
        expected_idx = torch.tensor([[2, 3, 4]], dtype=torch.int64)
        assert torch.equal(nn_idx, expected_idx)
        # multiple point test
        test_x = torch.tensor([1.6, 3], dtype=torch.float32)  # shape would be automatically expanded
        nn_idx = nn_util.find_nn_idx(test_x)
        nn_idx = torch.sort(nn_idx, dim=-1)[0]
        expected_idx = torch.tensor([[1, 2, 3], [2, 3, 4]], dtype=torch.int64)
        assert torch.equal(nn_idx, expected_idx)
        # with less k
        nn_idx = nn_util.find_nn_idx(test_x, k=1)
        expected_idx = torch.tensor([[2], [3]], dtype=torch.int64)
        assert torch.equal(nn_idx, expected_idx)
        # build sequential nn idx
        seq_nn_idx = nn_util.build_sequential_nn_idx(train_x)
        seq_nn_idx = torch.sort(seq_nn_idx, dim=-1)[0]
        expected_idx = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.int64)
        assert torch.equal(expected_idx, seq_nn_idx)
        # reset
        # set new train data
        train_x = torch.tensor([5.0, 3, 1, 4]).unsqueeze(-1)
        nn_util.set_nn_idx(train_x)
        test_x = torch.tensor([2.0, 10])
        nn_idx = nn_util.find_nn_idx(test_x)
        nn_idx = torch.sort(nn_idx, dim=-1)[0]
        expected_idx = torch.tensor([[1, 2, 3], [0, 1, 3]], dtype=torch.int64)
        assert torch.equal(nn_idx, expected_idx)
        seq_train_x = torch.tensor([10.0, 2, 5, 6, 9]).unsqueeze(-1)
        seq_nn_idx = nn_util.build_sequential_nn_idx(seq_train_x)
        seq_nn_idx = torch.sort(seq_nn_idx, dim=-1)[0]
        expected_idx = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int64)
        assert torch.equal(seq_nn_idx, expected_idx)