Fix model loading and device map inference

app.py

@@ -8,15 +8,14 @@ model_name = "ai4bharat/Airavata"
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Automatically determine the device map
-device_map = infer_auto_device_map(model_name)
-
-# Load the model with the device map
-model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
-    device_map=device_map, 
-    load_in_8bit=True  # Use 8-bit precision for reduced memory usage
-)
+# Load the model first
+model = AutoModelForCausalLM.from_pretrained(model_name, load_in_8bit=True)
+
+# Now infer the device map
+device_map = infer_auto_device_map(model)
+
+# Move model to the appropriate device based on device_map
+model.to(device_map)
 
 # Define the inference function
 def generate_text(prompt):