fix

app.py

@@ -56,11 +56,13 @@ def load_model(model_name):
         from skeleton_modeling_bert import SkeletonBertForMaskedLM
         tokenizer = BertTokenizer.from_pretrained(model_name)
         model = BertForMaskedLM.from_pretrained(model_name)
+        skeleton_model = SkeletonBertForMaskedLM
     elif model_name.startswith('roberta'):
         from transformers import RobertaTokenizer, RobertaForMaskedLM
         from skeleton_modeling_roberta import SkeletonRobertaForMaskedLM
         tokenizer = RobertaTokenizer.from_pretrained(model_name)
         model = RobertaForMaskedLM.from_pretrained(model_name)
+        skeleton_model = SkeletonRobertaForMaskedLM
     return tokenizer,model,skeleton_model
 
 def clear_data():