Mrw33554432 commited on
Commit
9d04c92
1 Parent(s): a4f5f9e

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -52,7 +52,7 @@ def quick_test(model, tokenizer, prompt: str):
52
  torch.set_default_device("cuda")
53
 
54
  tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5", trust_remote_code=True)
55
- model = AutoModelForCausalLM.from_pretrained("Mrw33554432/bitLinear-phi-1.5", trust_remote_code=True)
56
 
57
  print(model)
58
  # Replace Linear layers with BitLinear
 
52
  torch.set_default_device("cuda")
53
 
54
  tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5", trust_remote_code=True)
55
+ model = AutoModelForCausalLM.from_pretrained("Mrw33554432/bitLinear-phi-1.5", trust_remote_code=True, torch_dtype=torch.float16)
56
 
57
  print(model)
58
  # Replace Linear layers with BitLinear