fix(server): t5 cannot run in f16 (#356)

Fix #349
2023-05-23 12:15:54 +02:00 · 2023-05-23 12:15:54 +02:00 · 4f4c9c1665
parent 91d9beec90
commit 4f4c9c1665
1 changed files with 1 additions and 1 deletions
--- a/server/text_generation_server/models/t5.py
+++ b/server/text_generation_server/models/t5.py
@ -40,7 +40,7 @@ class T5Sharded(Seq2SeqLM):
        self.process_group, rank, world_size = initialize_torch_distributed()
        if torch.cuda.is_available():
            device = torch.device(f"cuda:{rank}")
-            dtype = torch.float16
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
        else:
            device = torch.device("cpu")
            dtype = torch.float32