gregH commited on
Commit
95469bd
1 Parent(s): 4a42d87

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -43,12 +43,12 @@ print(f"Starting to load the model to memory")
43
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
44
 
45
  m = AutoModelForCausalLM.from_pretrained(
46
- "google/gemma-2b-it", torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, trust_remote_code=True,token=os.environ["HF_Token"])
47
  embedding_func=m.get_input_embeddings()
48
  embedding_func.weight.requires_grad=False
49
  m = m.to(device)
50
 
51
- tok = AutoTokenizer.from_pretrained("google/gemma-2b-it", trust_remote_code=True,token=os.environ["HF_Token"])
52
  tok.padding_side = "left"
53
  tok.pad_token_id = tok.eos_token_id
54
  # using CUDA for an optimal experience
 
43
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
44
 
45
  m = AutoModelForCausalLM.from_pretrained(
46
+ "google/gemma-2b-it", torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, trust_remote_code=True)
47
  embedding_func=m.get_input_embeddings()
48
  embedding_func.weight.requires_grad=False
49
  m = m.to(device)
50
 
51
+ tok = AutoTokenizer.from_pretrained("google/gemma-2b-it", trust_remote_code=True)
52
  tok.padding_side = "left"
53
  tok.pad_token_id = tok.eos_token_id
54
  # using CUDA for an optimal experience