search

octo-models / octo

Octo is a transformer-based robot policy trained on a diverse mix of 800k robot trajectories.
https://octo-models.github.io/
MIT License
787 stars 152 forks source link

Memory allocation on GPU #109

Closed AasherH closed 3 months ago

AasherH commented 3 months ago

I am loading the octo-small model on my GPU (NVIDIA GeForce RTX 4090) and checking nvidia-smi shows that it uses about 20GB of memory storage, which seems high to me. Upon loading the octo-base model, it still takes up about 20GB of storage, which doesn't make much sense. Does anyone know how I can decrease the amount of memory allocated to loading the pretrained octo model, or do these numbers sound about right to everyone?

ojh6404 commented 3 months ago

See https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html

AasherH commented 3 months ago

Exactly what I needed, thank you!