factored_attention.py throws AssertionError

[agilebean]

After changing the artist and lyrics,
the jupyter notebook after line:
zs = _sample(zs, labels, sampling_kwargs, [None, None, top_prior], [2], hps)
throws the following error:

Sampling level 2
Sampling 8192 tokens for [0,8192]. Conditioning on 0 tokens
Ancestral sampling 3 samples with temp=0.98, top_k=0, top_p=0.0
0/8192 [00:00<?, ?it/s]

---------------------------------------------------------------------------

AssertionError                            Traceback (most recent call last)

<ipython-input-44-c3213f092598> in <module>()
      1 zs = [t.zeros(hps.n_samples,0,dtype=t.long, device='cuda') for _ in range(len(priors))]
----> 2 zs = _sample(zs, labels, sampling_kwargs, [None, None, top_prior], [2], hps)

6 frames

/usr/local/lib/python3.6/dist-packages/jukebox/transformer/factored_attention.py in check_cache(self, n_samples, sample_t, fp16)
    411 
    412     def check_cache(self, n_samples, sample_t, fp16):
--> 413         assert self.sample_t == sample_t, f"{self.sample_t} != {sample_t}"
    414         if sample_t == 0:
    415             assert self.cache == {}

AssertionError: 3344 != 0

[heewooj]

Hm.. this seems very strange. Could you share how to reproduce this error?

[agilebean]

I just changed the genre to Pop, and the artist to Katy Perry.
The text I entered was shorter, but could that cause an error?

[heewooj]

No, we have functions that pad short lyrics. I think top-level sampling was probably interrupted before which explains why the transformer has cached states already. Restarting the notebook kernel will most likely fix this.