Update README.md

README.md

@@ -127,7 +127,7 @@ from datasets import load_dataset
 
 # config
 model_id = "kotoba-tech/kotoba-whisper-v1.0"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
 generate_kwargs = {"language": "japanese", "task": "transcribe"}
@@ -189,7 +189,7 @@ from datasets import load_dataset
 
 # config
 model_id = "kotoba-tech/kotoba-whisper-v1.0"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
 generate_kwargs = {"language": "japanese", "task": "transcribe"}
@@ -225,7 +225,7 @@ from datasets import load_dataset
 
 # config
 model_id = "kotoba-tech/kotoba-whisper-v1.0"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
 generate_kwargs = {"language": "japanese", "task": "transcribe"}
@@ -303,7 +303,7 @@ from evaluate import load
 
 # model config
 model_id = "kotoba-tech/kotoba-whisper-v1.0"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
 generate_kwargs = {"language": "japanese", "task": "transcribe"}