add torch autocast to modeling

modeling.py

@@ -141,8 +141,9 @@ class MBZTestModelForCausalLM(PreTrainedModel):
         self.out_head = nn.Linear(d_model, n_vocab)
 
     def forward(self, x):
-        x = self.embed(x)
-        for block in self.blocks:
-            x = block(x)
-        x = self.out_head(self.norm(x))
-        return CausalLMOutput(logits=x)
+        with torch.autocast('cuda', dtype=torch.bfloat16):
+            x = self.embed(x)
+            for block in self.blocks:
+                x = block(x)
+            x = self.out_head(self.norm(x))
+            return CausalLMOutput(logits=x)