Update utils.py

utils.py

@@ -10,6 +10,7 @@ import requests
 import re
 import html
 import torch 
+from torch import cuda, bfloat16
 import sys
 import gc
 from pygments.lexers import guess_lexer, ClassNotFound
@@ -18,7 +19,7 @@ from pygments import highlight
 from pygments.lexers import guess_lexer,get_lexer_by_name
 from pygments.formatters import HtmlFormatter
 import transformers
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig,
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import datasets
 from datasets import load_dataset
 import evaluate
@@ -335,12 +336,11 @@ def daten_laden(name):
 
 
 #Quantisation - tzo speed up training
-def bnb_config (load4Bit, double_quant)
-    compute_dtype = getattr(torch, "float16")
+def bnb_config (load4Bit, double_quant):
     bnb_config = BitsAndBytesConfig(
         load_in_4bit= load4Bit,
         bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=compute_dtype,
+        bnb_4bit_compute_dtype=bfloat16,
         bnb_4bit_use_double_quant=double_quant,
     )
 return bnb_config