Update app.py

app.py

@@ -66,7 +66,7 @@ def make_layers(cfg, in_channels = 3,batch_norm=False,dilation = False):
 
 
 # Load the CSRNet model
-csrmodel = CSRNet()
+csrmodel = CSRNet(load_weights=True).cpu()
 checkpoint = torch.load("model.pt")
 csrmodel.load_state_dict(checkpoint)
 csrmodel.eval()
@@ -82,7 +82,7 @@ transform = transforms.Compose([
 # Define the prediction function
 def predict_count(input_image):
     # Preprocess the input image
-    image = transform(input_image).unsqueeze(0)
+    image = transform(input_image).unsqueeze(0).cpu()
 
     # Perform the forward pass
     output = csrmodel(image)