UPD8: clean-up and it runs now, cba'd to actually test it demo below. Fixes are from gh:madebyollin, check out his gist which actually works
weights now with ready to use checkpoint (will be up in few minutes)
Confirmed working with this code:
import time