Upload app.py

app.py

@@ -5,9 +5,9 @@ from clarifai_grpc.grpc.api.status import status_code_pb2
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-#torch.set_default_device("cpu")
-
-model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True)
+torch.set_default_tensor_type(torch.FloatTensor)  # Set the default tensor type to float32
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True).to(device)
 tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
 
 # GPT-4 credentials