Setting torch_dtype to bfloat16

CarperAI · Jul 22, 2023 · c40cdc4 · c40cdc4
1 parent a21a47a
commit c40cdc4
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/ppo_sentiments_8bit.py b/examples/ppo_sentiments_8bit.py
@@ -39,6 +39,7 @@ def main(hparams={}):
     # Set the model loading in 8 bits
     config.model.from_pretrained_kwargs = {
         "load_in_8bit": True,
+        "torch_dtype": torch.bfloat16,
         "device_map": "auto",
     }