about 'An Easier Trick for Boundary Prediction'

In your paper, we can get the predicted boundary as follows:

then I implemented 'An Easier Trick for Boundary Prediction' in my repo following the trick:
https://github.com/keonlee9420/DiffSinger/blob/f849f8def5abb38ad272a384e8bec838ea1957a4/boundary_predictor.py#L14-L45

and there are some helper functions for that (please focus on expected_kld_t and expected_kld_T function):
https://github.com/keonlee9420/DiffSinger/blob/f849f8def5abb38ad272a384e8bec838ea1957a4/model/diffusion.py#L351-L389

But as I noted in my README.md (in 2. of note section), the predicted boundary of LJSpeech is 100, which is the same as the total timesteps in Naive version.

So I'd like to ask you to briefly check my implementation. Could you please take a look at it and let me know if I missed something? Why do you think my boundary predictor shows unexpected K_step?

FYI, here is the sample output log of running boundary_predictor.py:

==================================== Prediction Configuration ====================================
 ---> Total Batch Size: 48
 ---> Path of ckpt: ./output/ckpt/LJSpeech_shallow_el_4
================================================================================================
100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 11/11 [00:08<00:00,  1.34it/s]
[tensor(6959.2134, device='cuda:0'), tensor(933.3702, device='cuda:0'), tensor(403.9860, device='cuda:0'), tensor(249.2317, device='cuda:0'), tensor(183.4001, device='cuda:0'), tensor(149.2621, device='cuda:0'), tensor(129.2204, device='cuda:0'), tensor(116.2622, device='cuda:0'), tensor(107.4923, device='cuda:0'), tensor(101.0867, device='cuda:0'), tensor(96.2093, device='cuda:0'), tensor(92.4524, device='cuda:0'), tensor(89.3728, device='cuda:0'), tensor(86.7645, device='cuda:0'), tensor(84.4990, device='cuda:0'), tensor(82.5240, device='cuda:0'), tensor(80.7848, device='cuda:0'), tensor(79.1111, device='cuda:0'), tensor(77.5320, device='cuda:0'), tensor(76.0396, device='cuda:0'), tensor(74.6199, device='cuda:0'), tensor(73.2726, device='cuda:0'), tensor(71.9328, device='cuda:0'), tensor(70.6272, device='cuda:0'), tensor(69.2854, device='cuda:0'), tensor(68.0120, device='cuda:0'), tensor(66.7351, device='cuda:0'), tensor(65.4260, device='cuda:0'), tensor(64.1837, device='cuda:0'), tensor(62.9117, device='cuda:0'), tensor(61.6452, device='cuda:0'), tensor(60.3592, device='cuda:0'), tensor(59.0823, device='cuda:0'), tensor(57.8210, device='cuda:0'), tensor(56.5481, device='cuda:0'), tensor(55.2716, device='cuda:0'), tensor(54.0141, device='cuda:0'), tensor(52.7686, device='cuda:0'), tensor(51.4833, device='cuda:0'), tensor(50.2068, device='cuda:0'), tensor(48.9261, device='cuda:0'), tensor(47.6881, device='cuda:0'), tensor(46.4407, device='cuda:0'), tensor(45.2071, device='cuda:0'), tensor(43.9496, device='cuda:0'), tensor(42.7181, device='cuda:0'), tensor(41.5266, device='cuda:0'), tensor(40.2994, device='cuda:0'), tensor(39.1266, device='cuda:0'), tensor(37.9398, device='cuda:0'), tensor(36.7822, device='cuda:0'), tensor(35.6130, device='cuda:0'), tensor(34.5006, device='cuda:0'), tensor(33.3484, device='cuda:0'), tensor(32.2580, device='cuda:0'), tensor(31.1593, device='cuda:0'), tensor(30.1051, device='cuda:0'), tensor(29.0614, device='cuda:0'), tensor(28.0244, device='cuda:0'), tensor(27.0115, device='cuda:0'), tensor(26.0248, device='cuda:0'), tensor(25.0589, device='cuda:0'), tensor(24.1051, device='cuda:0'), tensor(23.1736, device='cuda:0'), tensor(22.2743, device='cuda:0'), tensor(21.3856, device='cuda:0'), tensor(20.5282, device='cuda:0'), tensor(19.6825, device='cuda:0'), tensor(18.8733, device='cuda:0'), tensor(18.0839, device='cuda:0'), tensor(17.3134, device='cuda:0'), tensor(16.5815, device='cuda:0'), tensor(15.8417, device='cuda:0'), tensor(15.1426, device='cuda:0'), tensor(14.4522, device='cuda:0'), tensor(13.8025, device='cuda:0'), tensor(13.1645, device='cuda:0'), tensor(12.5432, device='cuda:0'), tensor(11.9491, device='cuda:0'), tensor(11.3789, device='cuda:0'), tensor(10.8328, device='cuda:0'), tensor(10.2960, device='cuda:0'), tensor(9.7815, device='cuda:0'), tensor(9.2841, device='cuda:0'), tensor(8.8136, device='cuda:0'), tensor(8.3660, device='cuda:0'), tensor(7.9211, device='cuda:0'), tensor(7.5027, device='cuda:0'), tensor(7.1040, device='cuda:0'), tensor(6.7245, device='cuda:0'), tensor(6.3511, device='cuda:0'), tensor(6.0048, device='cuda:0'), tensor(5.6679, device='cuda:0'), tensor(5.3475, device='cuda:0'), tensor(5.0427, device='cuda:0'), tensor(4.7507, device='cuda:0'), tensor(4.4784, device='cuda:0'), tensor(4.2143, device='cuda:0'), tensor(3.9639, device='cuda:0'), tensor(3.7258, device='cuda:0')]
tensor(0.2382, device='cuda:0')

Predicted Boundary K is 100

Thanks in advance!

Start with boundary_predictor.py and the expected_kld_t and expected_kld_T functions in model/diffusion.py, then read note 2 in README.md alongside the reported output. Trace how the predicted boundary K is produced and compare it with the described trick. Done means explaining why K_step is 100 and identifying whether the implementation misses anything.