update device

app.py

@@ -92,8 +92,8 @@ def main():
     die_model = UNetDIEModel(args=args)
 
     # Partially apply the model and device arguments to die_inference
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    partial_die_inference = partial(die_inference, device=device, die_model=die_model)
+    args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    partial_die_inference = partial(die_inference, device=args.device, die_model=die_model)
 
     # Gradio Interface
     with gr.Blocks() as demo: