tying to solve RuntimeError: CUDA must not be initialized in the main process on Spaces with Stateless GPU environment.

app.py

@@ -75,6 +75,8 @@ def download_smplx():
 def infer(fn, det_thresh, nms_kernel_size):
     global device
     global model
+
+    model = model.to(device)
     
     # Is it an image from example_data_dir ?
     basename = Path(os.path.basename(fn)).stem
@@ -215,7 +217,8 @@ if __name__ == "__main__":
         download_smplx()
 
         # Loading the model
-        model = _load_model(args.model_name, device=device)
+        # model = _load_model(args.model_name, device=device)
+        model = _load_model(args.model_name, device=torch.device('cpu'))
         
         # Gradio demo
         with gr.Blocks(title="Multi-HMR", css=".gradio-container") as demo: