add pytorch cuda

Dockerfile

@@ -0,0 +1,25 @@
+FROM python:3.9
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+ PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+CMD ["python", "main.py"]

app.py

@@ -5,6 +5,7 @@ import gradio as gr
 model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3")
 
 if torch.cuda.is_available():
+ print("Using GPU")
  model = model.cuda()
 
 convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")

requirements.txt

@@ -1,4 +1,5 @@
 av
+--extra-index-url https://download.pytorch.org/whl/cu113
 torch==2.1.0
 torchaudio==2.1.0
 torchvision==0.16.0