bankholdup commited on
Commit
4d0a9ee
1 Parent(s): f5b970b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -10,7 +10,8 @@ def load_model():
10
  # model_ckpt = "bankholdup/rugpt3_song_writer"
11
  model_ckpt = "bankholdup/mgpt_song_writer"
12
  tokenizer = GPT2Tokenizer.from_pretrained(model_ckpt)
13
- model = GPT2LMHeadModel.from_pretrained(model_ckpt)
 
14
  return tokenizer, model
15
 
16
  def set_seed(rng=100000):
 
10
  # model_ckpt = "bankholdup/rugpt3_song_writer"
11
  model_ckpt = "bankholdup/mgpt_song_writer"
12
  tokenizer = GPT2Tokenizer.from_pretrained(model_ckpt)
13
+ # model = GPT2LMHeadModel.from_pretrained(model_ckpt)
14
+ model = GPT2LMHeadModel.from_pretrained(model_ckpt, revision="float16", torch_dtype=torch.float16, low_cpu_mem_usage=True)
15
  return tokenizer, model
16
 
17
  def set_seed(rng=100000):