Pytorch NLP Huggingface: model not loaded on GPU
I have this code that init a class with a model and a tokenizer from Huggingface.
On Google Colab this code works fine, it loads the model on the GPU memory without problems.
On Google Cloud Platform it does not work, it loads the model on gpu, whatever I try.


class OPT:
    def __init__(self, model_name: str = "facebook/opt-2.7b", use_gpu: bool = False):
        self.model_name = model_name
        self.use_gpu = use_gpu and torch.cuda.is_available()
        print(f"Use gpu:: {self.use_gpu}")

        if self.use_gpu:
            print("Using gpu")
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name, torch_dtype=torch.float16
            ).cuda()
        else:
            print("Using cpu")
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name, torch_dtype=torch.float32, low_cpu_mem_usage=True
            )

        # the fast tokenizer currently does not work correctly
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)



The printed output is correct:


Use gpu:: True
Using gpu



But the nvidia-smi says that there is no process running on the gpu:


+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.82.01    Driver Version: 470.82.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:04.0 Off |                    0 |
| N/A   40C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+



And with htop I can see that the process is using the cpu ram.

You should use the .to(device) method like this:
device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)
model = nameofyourmodel.to(device)

