Add missing elements to sft_trainer document (#1029)

huggingface · Nov 23, 2023 · 3719f7a · 3719f7a
1 parent e7961e4
commit 3719f7a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/docs/source/sft_trainer.mdx b/docs/source/sft_trainer.mdx
@@ -232,6 +232,7 @@ You can directly pass the kwargs of the `from_pretrained()` method to the [`SFTT
 
 ```python
 model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", torch_dtype=torch.bfloat16)
+```
 
 ```python
 ...
@@ -240,6 +241,9 @@ trainer = SFTTrainer(
     "facebook/opt-350m",
     train_dataset=dataset,
     dataset_text_field="text",
+    model_init_kwargs={
+        "torch_dtype": torch.bfloat16,
+    },
 )
 
 trainer.train()