google-research / big_transfer

Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.
https://arxiv.org/abs/1912.11370
Apache License 2.0
1.5k stars 175 forks source link

Model taking entire GPU #28

Closed sumedhvdatar closed 4 years ago

sumedhvdatar commented 4 years ago

BIT model is lite weight but the model is taking entire GPU memory. Need a way to free up memory.

lucasb-eyer commented 4 years ago

This is not nearly enough information. Some frameworks, for example TF and Jax, allocate all GPU memory by default. This is a framework-level setting and not related to our code release.