Support bfloat16 data type (#54)

2023-05-03 14:09:44 -07:00
parent 436e523bf1
commit e070829ae8
12 changed files with 455 additions and 53 deletions
--- a/cacheflow/models/utils.py
+++ b/cacheflow/models/utils.py
@@ -17,6 +17,7 @@ _STR_DTYPE_TO_TORCH_DTYPE = {
    'float': torch.float,
    'float16': torch.float16,
    'float32': torch.float32,
+    'bfloat16': torch.bfloat16,
 }