Update app.py

app.py

@@ -5,9 +5,11 @@ import torch
 
 os.environ['CURL_CA_BUNDLE'] = ''
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 config = LinearMappingConfig()
 model = LinearMapping(config)
-model.load_state_dict(torch.load("pytorch_model.bin"))
+model.load_state_dict(torch.load("pytorch_model.bin", map_location=device))
 processor = LinearMappingProcessor(config)
 processor.tokenizer.padding_side = 'left'
 processor.tokenizer.pad_token_id = processor.tokenizer.eos_token_id
@@ -53,4 +55,4 @@ demo = gr.Interface(
     description=description
 )
 
-demo.launch(share=True)
+demo.launch()