send model to device

app.py

@@ -136,7 +136,7 @@ if __name__ == "__main__":
     lora = False
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     model = AudioTokenWrapper(lora, device)
-
+    model = model.to(device)
     description = """<p>
     This is a demo of <a href='https://pages.cs.huji.ac.il/adiyoss-lab/AudioToken' target='_blank'>AudioToken: Adaptation of Text-Conditioned Diffusion Models for Audio-to-Image Generation</a>.<br><br>
     A novel method utilizing latent diffusion models trained for text-to-image-generation to generate images conditioned on audio recordings. Using a pre-trained audio encoding model, the proposed method encodes audio into a new token, which can be considered as an adaptation layer between the audio and text representations.<br><br>