Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 import torch.nn as nn
 from torchvision import models, transforms
 from PIL import Image
-import os
 
 # Define the same custom residual block and EfficientNetWithNovelty model
 class ResidualBlock(nn.Module):
@@ -53,11 +52,11 @@ class EfficientNetWithNovelty(nn.Module):
         
         return x
 
-# Load the model checkpoint
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Load the model checkpoint on CPU
+device = torch.device('cpu')  # Ensure it's using CPU
 num_classes = 10  # Number of classes as per your dataset
 model = EfficientNetWithNovelty(num_classes)
-checkpoint = torch.load('best_model2.pth')
+checkpoint = torch.load('best_model2.pth', map_location=torch.device('cpu'))
 model.load_state_dict(checkpoint['model_state_dict'])
 model.to(device)
 model.eval()