update model code

app.py

@@ -299,7 +299,7 @@ def predict():
 
     def remove_module_from_checkpoint(checkpoint):
         new_state_dict = {}
-        for key, value in checkpoint["model_state_dict"].items():
+        for key, value in checkpoint.items():
             new_key = key.replace("module.", "")
             new_state_dict[new_key] = value
         checkpoint["model_state_dict"] = new_state_dict