Update StreamlitModel.py

StreamlitModel.py

@@ -20,8 +20,8 @@ import torch.hub
 profanity.load_censor_words()    
 
 #It seems the streamlit space does not allow mps
-#device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_built() else 'cpu'
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_built() else 'cpu'
+#device = 'cuda' if torch.cuda.is_available() else 'cpu'
     
 class GPT2_Model(GPT2PreTrainedModel):
 
@@ -105,7 +105,7 @@ def TextGeneration(prompt,prof=False,parts=True):
     gpt_model = GPT2_Model(configuration).to(device)
     #gpt_model.load_state_dict(torch.load('GPT-Trained-Model-Prod.pt'))
     state_dict = torch.hub.load_state_dict_from_url(r'https://github.com/Selbl/LyricGeneration/raw/main/GPT-Trained-Model-Prod.pt?download=')
-    gpt_model.load_state_dict(state_dict)
+    gpt_model.load_state_dict(state_dict,map_location=device)
     #Load tokenizer
     tokenizer = GPT2Tokenizer.from_pretrained('gpt2', pad_token='<|pad|>') 
     #Call loader function