"Cannot allocate memory" error whan running locally with smaller model

[tomfpritchard]

Hi, great work here. I tried the colab and I could not be more impressed.

I'm getting this memory error when I try to use either model. Machine has a 1080ti and 16GB RAM.

I've tried reducing the max_batch_size in sample.py and the number of samples in the call, but nothing seems to change.

Downloading from gce
Restored from /home/tom/.cache/jukebox-assets/models/5b/prior_level_0.pth.tar
0: Loading prior in eval mode
Conditioning on 1 above level(s)
Checkpointing convs
Checkpointing convs
Loading artist IDs from /home/tom/jukebox/jukebox/data/ids/v2_artist_ids.txt
Loading artist IDs from /home/tom/jukebox/jukebox/data/ids/v2_genre_ids.txt
Level:1, Cond downsample:4, Raw to tokens:32, Sample length:262144
Downloading from gce
Traceback (most recent call last):
File "jukebox/sample.py", line 278, in
fire.Fire(run)
File "/home/tom/anaconda3/envs/jukebox/lib/python3.7/site-packages/fire/core.py", line 127, in Fire
component_trace = _Fire(component, args, context, name)
File "/home/tom/anaconda3/envs/jukebox/lib/python3.7/site-packages/fire/core.py", line 366, in _Fire
component, remaining_args)
File "/home/tom/anaconda3/envs/jukebox/lib/python3.7/site-packages/fire/core.py", line 542, in _CallCallable
result = fn(*varargs, **kwargs)
File "jukebox/sample.py", line 275, in run
save_samples(model, device, hps, sample_hps)
File "jukebox/sample.py", line 180, in save_samples
vqvae, priors = make_model(model, device, hps)
File "/home/tom/jukebox/jukebox/make_models.py", line 185, in make_model
priors = [make_prior(setup_hparams(priors[level], dict()), vqvae, 'cpu') for level in levels]
File "/home/tom/jukebox/jukebox/make_models.py", line 185, in
priors = [make_prior(setup_hparams(priors[level], dict()), vqvae, 'cpu') for level in levels]
File "/home/tom/jukebox/jukebox/make_models.py", line 169, in make_prior
restore(hps, prior, hps.restore_prior)
File "/home/tom/jukebox/jukebox/make_models.py", line 54, in restore
checkpoint = load_checkpoint(checkpoint_path)
File "/home/tom/jukebox/jukebox/make_models.py", line 34, in load_checkpoint
download(gs_path, local_path)
File "/home/tom/jukebox/jukebox/utils/gcs_utils.py", line 36, in download
subprocess.call(args)
File "/home/tom/anaconda3/envs/jukebox/lib/python3.7/subprocess.py", line 339, in call
with Popen(*popenargs, **kwargs) as p:
File "/home/tom/anaconda3/envs/jukebox/lib/python3.7/subprocess.py", line 800, in init
restore_signals, start_new_session)
File "/home/tom/anaconda3/envs/jukebox/lib/python3.7/subprocess.py", line 1482, in _execute_child
restore_signals, start_new_session, preexec_fn)
OSError: [Errno 12] Cannot allocate memory