@koustuvsinha have you figured out if this implementation is correct? I am also running into the same issue as you while attempting to implement the schedules sampling method. If you’ve figured this out, do you have any good resources you refer to?