Update app.py

app.py

@@ -13,6 +13,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 os.environ['HF_HOME'] = '/app/.cache'
 os.environ['XDG_CACHE_HOME'] = '/app/.cache'
 
+
 app = Flask(__name__)
 CORS(app)
 
@@ -21,32 +22,32 @@ model_loaded = False
 load_error = None
 generator = None
 
-# --------------------------------------------------
-# Asynchronous Model Loading
-# --------------------------------------------------
 def load_model():
     global model_loaded, load_error, generator
     try:
-        # Initialize model with low-memory settings
+        # Detect device and dtype automatically
+        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        
         model = AutoModelForCausalLM.from_pretrained(
             "gpt2-medium",
             use_safetensors=True,
             device_map="auto",
-            low_cpu_mem_usage=True,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True
         )
         
         tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
         
+        # Initialize pipeline without explicit device assignment
         generator = pipeline(
             'text-generation',
             model=model,
             tokenizer=tokenizer,
-            device=0 if torch.cuda.is_available() else -1
+            torch_dtype=dtype
         )
         
         model_loaded = True
-        print("Model loaded successfully")
+        print(f"Model loaded on {model.device}")
         
     except Exception as e:
         load_error = str(e)
@@ -55,6 +56,8 @@ def load_model():
 # Start model loading in background thread
 Thread(target=load_model).start()
 
+
+
 # --------------------------------------------------
 # IEEE Format Template
 # --------------------------------------------------