refactor

app.py

@@ -6,7 +6,10 @@ import transformers
 from transformers import AutoTokenizer
 
 logging.basicConfig(level=logging.INFO)
-logging.info(f"APP startup")
+if torch.cuda.is_available():
+    logging.info("Running on GPU")
+else:
+    logging.info("Running on CPU")
 
 if "googleflan" == os.environ.get("MODEL"):
     model = "google/flan-t5-small"
@@ -22,8 +25,8 @@ elif "llama" == os.environ.get("MODEL"):
     pipeline = transformers.pipeline(
         "text-generation",
         model=model,
-        # torch_dtype=torch.float32,
-        torch_dtype="auto",
+        torch_dtype=torch.float16,
+        # torch_dtype="auto",
         low_cpu_mem_usage=True,
         device_map="auto",
         token=os.environ.get("HF_TOKEN"),