Turn autocast back on

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)