Update app.py

app.py

@@ -9,7 +9,7 @@ import torch.nn as nn
 import timm 
 
 REPO_ID = "Raaniel/model-smoke"
-MODEL_FILE_NAME = "model_smoke.pt"
+MODEL_FILE_NAME = "best_model_epoch_32.pth"
 USE_CUDA = torch.cuda.is_available()
 num_classes = 3
 
@@ -32,7 +32,7 @@ model.classifier = nn.Sequential(
 )
 
 # Load the model weights
-model.load_state_dict(state['weights'])
+model.load_state_dict(state)
 
 # Move model to the appropriate device
 device = torch.device('cuda' if USE_CUDA else 'cpu')