Unable to run on GPU

[Ntweat]

I am trying to run Nerfies on GPU it always "RuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16504591536 bytes."

I tried it locally and on colab (selected GPU runtime and GPU in code)

Both places same error.

It happens at cell "Train a Nerfie!"
on line "with time_tracker.record_time('train_step'):
state, stats, keys = ptrain_step(keys, state, batch, scalar_params)
time_tracker.toc('total')"

Any solution to this??

[keunhong]

Hi, unfortuantely since Colab GPUs are a bit limited you might have to reduce the number of samples, number of training images, or network size to get it to fit on those GPUs.

I had better luck with the Cloud TPUs so I suggested trying that. Otherwise the only solution is to run it on your own GPU machine.

[keunhong]

Please feel free to re-open if you encounter any other issues.