Update mymodule/utils.py

mymodule/utils.py

@@ -1,5 +1,6 @@
 import os
 
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
@@ -13,5 +14,5 @@ def load_model_and_tok():
         "meta-llama/Llama-3.2-3B-Instruct",
         token=AUTH_TOKEN,
     )
-    model = model.to("cuda").eval()
+    model = model.to("cuda" if torch.cuda.is_available() else "cpu").eval()
     return model, tok