Print JAX device

app.py

@@ -5,6 +5,7 @@ import os
 import gradio as gr
 from PIL import Image
 import numpy as np
+import jax
 
 from gradio_dualvision import DualVisionApp
 from gradio_dualvision.gradio_patches.radio import Radio
@@ -14,6 +15,9 @@ from super_resolve import process
 
 REPO_ID = "prs-eth/thera-edsr-plus"
 
+print(f"JAX devices: {jax.devices()}")
+print(f"JAX device type: {jax.devices()[0].device_kind}")
+
 # load model
 model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pkl")
 with open(model_path, 'rb') as fh: