kpriyanshu256 commited on
Commit
4013271
1 Parent(s): 764b63c

Memory release

Browse files
Files changed (2) hide show
  1. app.py +2 -1
  2. model.py +1 -1
app.py CHANGED
@@ -15,8 +15,9 @@ MAX_LEN = 256
15
  MODEL = BertAD()
16
  vec = MODEL.state_dict()['bert.embeddings.position_ids']
17
  chkp = torch.load(os.path.join('model', 'model_0.bin'), map_location='cpu')
18
- chkp['bert.embeddings.position_ids'] =vec
19
  MODEL.load_state_dict(chkp)
 
20
 
21
 
22
  def sample_text(text, acronym, max_len):
 
15
  MODEL = BertAD()
16
  vec = MODEL.state_dict()['bert.embeddings.position_ids']
17
  chkp = torch.load(os.path.join('model', 'model_0.bin'), map_location='cpu')
18
+ chkp['bert.embeddings.position_ids'] = vec
19
  MODEL.load_state_dict(chkp)
20
+ del chkp, vec
21
 
22
 
23
  def sample_text(text, acronym, max_len):
model.py CHANGED
@@ -6,7 +6,7 @@ class BertAD(nn.Module):
6
  def __init__(self):
7
  super(BertAD, self).__init__()
8
 
9
- model_config = transformers.AutoConfig.from_pretrained('model')
10
  model_config.update({"output_hidden_states":True})
11
 
12
  self.bert = transformers.BertModel(model_config)
 
6
  def __init__(self):
7
  super(BertAD, self).__init__()
8
 
9
+ model_config = transformers.AutoConfig.from_pretrained('./model')
10
  model_config.update({"output_hidden_states":True})
11
 
12
  self.bert = transformers.BertModel(model_config)