Update app.py

app.py

@@ -56,7 +56,7 @@ class Gradio:
 
         if os.path.isfile(model_path):
             model.load_state_dict(
-                torch.load(model_path)["model_state_dict"], strict=False
+                torch.load(model_path), strict=False, map_location=self.device
             )
 
         return model