BifrostTitan commited on
Commit
07b385f
·
verified ·
1 Parent(s): 2311fd0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -2
app.py CHANGED
@@ -65,6 +65,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
65
  hf_save_dir = "./bitnet"
66
  model = AutoModelForCausalLM.from_pretrained(
67
  model_id,
 
68
  device_map="auto"
69
  )
70
  device = model.device
@@ -73,8 +74,6 @@ if torch.cuda.is_available():
73
  else:
74
  print("CUDA not available; using CPU.")
75
 
76
-
77
- model.to(device)
78
  print(f"Loaded pre-trained Hugging Face model '{model_id}'.")
79
 
80
  # ---------------------------------------------------------------------------------
 
65
  hf_save_dir = "./bitnet"
66
  model = AutoModelForCausalLM.from_pretrained(
67
  model_id,
68
+ torch_dtype=torch.bfloat32
69
  device_map="auto"
70
  )
71
  device = model.device
 
74
  else:
75
  print("CUDA not available; using CPU.")
76
 
 
 
77
  print(f"Loaded pre-trained Hugging Face model '{model_id}'.")
78
 
79
  # ---------------------------------------------------------------------------------