slai-labs / get-beam

Run GPU inference and training jobs on serverless infrastructure that scales with you.

https://beam.cloud

Model not being moved to GPU

dhanushreddy291 opened this issue a year ago · comments

Dhanush Reddy commented a year ago

In this replit-code, I feel the model is not being moved to GPU:
https://github.com/slai-labs/get-beam/blob/main/examples/replit-code/replit-code.py

Also the cache_dir is not being used in the instantiation of AutoModelForCausalLM

Please correct me if I am wrong

Nolan Tremelling commented a year ago

Thanks for this catch! Fixed with #46

Dhanush Reddy commented a year ago

Cool 😎

btw, I think cache_dir can also be added to AutoTokenizer