use torch autocast for llama model generation

model.py

@@ -215,22 +215,24 @@ class SALMONN(nn.Module):
         embeds = torch.cat([bos_embeds, prompt_left_embeds, speech_embeds, prompt_right_embeds], dim=1)
         atts = torch.ones(embeds.size()[:-1], dtype=torch.long).to(embeds.device)
 
-        # generate
-        output = self.llama_model.generate(
-            inputs_embeds=embeds,
-            max_length=max_length,
-            num_beams=num_beams,
-            do_sample=do_sample,
-            min_length=min_length,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            length_penalty=length_penalty,
-            temperature=temperature,
-            attention_mask=atts,
-            bos_token_id=self.llama_tokenizer.bos_token_id,
-            eos_token_id=self.llama_tokenizer.eos_token_id,
-            pad_token_id=self.llama_tokenizer.pad_token_id
-        )
+        from torch.cuda.amp import autocast
+
+        with autocast(device_type="cuda", dtype=torch.float16):
+            output = self.llama_model.generate(
+                inputs_embeds=embeds,
+                max_length=max_length,
+                num_beams=num_beams,
+                do_sample=do_sample,
+                min_length=min_length,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                length_penalty=length_penalty,
+                temperature=temperature,
+                attention_mask=atts,
+                bos_token_id=self.llama_tokenizer.bos_token_id,
+                eos_token_id=self.llama_tokenizer.eos_token_id,
+                pad_token_id=self.llama_tokenizer.pad_token_id
+            )
         
         output_text = self.llama_tokenizer.batch_decode(output, add_special_tokens=False, skip_special_tokens=True)