Update app.py

app.py

@@ -4,9 +4,10 @@ import numpy as np
 import torch
 from torchvision import transforms
 from PIL import Image
+from TranSalNet_Res import TranSalNet
 
 # Load the model and set the device
-model = TranSalNet()  # Assuming you have defined your model
+model = TranSalNet()
 model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_location=torch.device('cpu')))
 model.eval()
 device = torch.device('cpu')