Remove autocast

llama_diffusion_model.py

@@ -137,14 +137,14 @@ class CustomTransformerModel(PreTrainedModel):
         attention_mask = attention_mask.to(dtype=torch.float32)  # required for SDPA and Flash attention
 
 
-        with autocast("cuda", dtype=torch.float16):
-            outputs = self.llama(
-                input_ids,
-                attention_mask=attention_mask,
-                output_hidden_states=True,
-                use_cache=False,
-                **kwargs
-            )
+        # with autocast("cuda", dtype=torch.float16):
+        outputs = self.llama(
+            input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            use_cache=False,
+            **kwargs
+        )
 
         logits = outputs.logits[:, :, :self.config.vocab_size].view(batch_size, seq_len, self.config.vocab_size)