CUDA or not CUDA

model_translation.py

@@ -25,7 +25,7 @@ model_names = {
 # Registry for all loaded bilingual models
 tokenizer_model_registry = {}
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'cpu'
 
 def get_tokenizer_model_for_src_lang(src_lang: str) -> (AutoTokenizer, AutoModelForSeq2SeqLM):
     """
@@ -47,7 +47,7 @@ def get_tokenizer_model_for_src_lang(src_lang: str) -> (AutoTokenizer, AutoModel
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     if model.config.torch_dtype != torch.float16:
         model = model.half()
-    model = model.to(device)
+    model.to(device)
     tokenizer_model_registry[src_lang] = (tokenizer, model)
 
     return (tokenizer, model)
@@ -65,5 +65,8 @@ model_MADLAD_name = "google/madlad400-3b-mt"
 #model_MADLAD_name = "google/madlad400-7b-mt-bt"
 tokenizer_multilingual = AutoTokenizer.from_pretrained(model_MADLAD_name, use_fast=True)
 model_multilingual = AutoModelForSeq2SeqLM.from_pretrained(
-    model_MADLAD_name, device_map="auto", torch_dtype=torch.float16)
+    model_MADLAD_name,
+    device_map="auto",
+    torch_dtype=torch.float16
+    low_cpu_mem_usage=True)