Update app.py

app.py

@@ -60,10 +60,9 @@ def load_model():
     model_name = 'model'
     cat_count = 358
 
-    checkpoint = torch.load(os.path.join(chkp_folder, f"{model_name}.pt"), weights_only=False)
+    checkpoint = torch.load(os.path.join(chkp_folder, f"{model_name}.pt"), weights_only=False, map_location=torch.device('cpu'))
 
     # Создаём те же классы, что и внутри чекпоинта
-    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
 
     model_ = AutoModelForSequenceClassification.from_pretrained('distilbert/distilbert-base-cased', num_labels=cat_count).to(device)
 
@@ -77,7 +76,7 @@ def load_model():
     model, optimizer = create_model_and_optimizer(model_)
 
     # Загружаем состояния из чекпоинта
-    model.load_state_dict(checkpoint['model_state_dict'], map_location=torch.device('cpu'))
+    model.load_state_dict(checkpoint['model_state_dict'])
     ind_to_cat = checkpoint['ind_to_cat']
     tokenizer = AutoTokenizer.from_pretrained('distilbert/distilbert-base-cased')
     return model, tokenizer, ind_to_cat