Fix dtype check for builder.py

Summary:
llama2 7b checkpoint is with bfloat16 hence this check doesnt work without including that.

Created from CodeHub with https://fburl.com/edit-in-codehub

Reviewed By: larryliu0820, manuelcandales

Differential Revision: D54122475

fbshipit-source-id: 4d25833776c9ca14f009550e18c6930734630586

Author: kimishpatel

examples/models/llama2/builder.py

@@ -75,7 +75,11 @@ def load_llama_model(
     )
     state_dict = model.state_dict()
     dtype = state_dict[next(iter(state_dict))].dtype
-    assert dtype in [torch.float16, torch.float32], "Only support fp16 or fp32"
+    assert dtype in [
+        torch.bfloat16,
+        torch.float16,
+        torch.float32,
+    ], f"Only support bfloat16, fp16 or fp32 got {dtype}"
     logging.info(f"Loaded model with dtype={dtype}")
 
     return LlamaEdgeManager(