Fix torchvision NMS registration on ZeroGPU

Dockerfile

@@ -3,7 +3,8 @@ FROM nvidia/cuda:12.1.0-devel-ubuntu22.04
 # Set environment variables
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
-    PYTHON_VERSION=3.10
+    PYTHON_VERSION=3.10 \
+    TORCHVISION_DISABLE_META_REGISTRATIONS=1
 
 # Install system dependencies
 RUN apt-get update && apt-get install -y \

app.py

@@ -10,6 +10,9 @@ import logging
 import warnings
 from pathlib import Path
 
+# Prevent torchvision from registering optional CUDA/Meta ops (nms) that may be missing on ZeroGPU
+os.environ.setdefault("TORCHVISION_DISABLE_META_REGISTRATIONS", "1")
+
 import gradio as gr
 import torch
 import numpy as np