Update model.py

model.py

@@ -238,7 +238,12 @@ tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_path'])
 path = 'model.pth'
 
 model = NBMEModel().to(config['device'])
-model.load_state_dict(torch.load(path, map_location=torch.device(config['device']))['model'])
+model.load_state_dict(
+    torch.load(
+        path, 
+        map_location=torch.device(config['device'])
+    )
+)
 model.eval()
 
 # input_text = create_sample_test()