update model code

app.py

@@ -297,7 +297,7 @@ def predict():
     checkpoint = torch.load(
         "/home/user/app/a.pth", weights_only=False, map_location=torch.device('cpu'))
     # checkpoint = remove_module_from_checkpoint(checkpoint)
-    model_new.load_state_dict(checkpoint['model_state_dict'])
+    model_new.load_state_dict(checkpoint)
     image = Image.open(temp_path).convert("RGB")
     output = model_new(transform(image).unsqueeze(0))
     stage = output.item()