Update app.py

app.py

@@ -17,11 +17,11 @@ def load_model():
             low_cpu_mem_usage=True      # Enable low memory usage on CPU
         )
     
-    # Move model to CPU
+    # Use 'auto' or 'balanced' device_map for CPU or GPU
     model = load_checkpoint_and_dispatch(
         model,
         MODEL_NAME,
-        device_map="cpu",  # Load model onto CPU
+        device_map="auto",  # Automatically determine the best device allocation
     )
     
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
@@ -57,4 +57,4 @@ if st.button("Generate Code"):
     else:
         st.warning("Please enter a prompt.")
 
-st.caption("Created by Shamil")
+st.caption("Powered by CodeGen-2B | Streamlit UI | CPU Optimized")