v1

trol/load_trol.py

@@ -64,8 +64,14 @@ def load_trol(link):
     # Loading tokenizer & Loading backbone model (error -> then delete flash attention)
     tok_trol = TroLTokenizer.from_pretrained(path, padding_side='left')
     try:
-        trol = TroLForCausalLM.from_pretrained(path, **huggingface_config).cuda()
+        trol = TroLForCausalLM.from_pretrained(path, **huggingface_config)
     except:
         del huggingface_config["attn_implementation"]
-        trol = TroLForCausalLM.from_pretrained(path, **huggingface_config).cuda()
+        trol = TroLForCausalLM.from_pretrained(path, **huggingface_config)
+
+    # wrapping
+    try:
+        trol = trol.cuda()
+    except:
+        pass
     return trol, tok_trol