Update app.py

app.py

@@ -45,7 +45,6 @@ def load_model():
     return model, model_config
     
 model, model_config = load_model()
-model = model.to(device,torch.bfloat16)
 
 # Function to set up, generate, and process the audio
 @spaces.GPU(duration=60)  # Allocate GPU only when this function is called
@@ -55,6 +54,8 @@ def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"Using device: {device}")
+    
+    model = model.to(device,torch.bfloat16)
 
     # Fetch the Hugging Face token from the environment variable
     hf_token = os.getenv('HF_TOKEN')