small change

app.py

@@ -89,10 +89,11 @@ def main():
         use_auth_token=die_token
     )
     
+    args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
     die_model = UNetDIEModel(args=args)
 
     # Partially apply the model and device arguments to die_inference
-    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