Update README.md

README.md

@@ -50,7 +50,7 @@ model_path = "Open-Bee/Bee-8B-RL"
 # Load model
 model = AutoModel.from_pretrained(
     model_path,
-    torch_dtype="float32",
+    dtype="bfloat16",
     trust_remote_code=True,
 ).to("cuda")
 

config.json

@@ -183,13 +183,13 @@
     "rope_scaling": null,
     "rope_theta": 1000000,
     "sliding_window": null,
-    "torch_dtype": "float32",
+    "torch_dtype": "bfloat16",
     "use_cache": true,
     "use_sliding_window": false,
     "vocab_size": 152000
   },
   "tie_word_embeddings": false,
-  "torch_dtype": "float32",
+  "torch_dtype": "bfloat16",
   "transformers_version": "4.55.0",
   "use_image_newline_parameter": true,
   "video_token_index": 151670,
@@ -209,7 +209,7 @@
     "num_channels": 3,
     "num_hidden_layers": 26,
     "patch_size": 14,
-    "torch_dtype": "float32",
+    "torch_dtype": "bfloat16",
     "vision_use_head": false
   },
   "vision_feature_layer": -1,