Dockerfile

# Use the official PyTorch Docker image as a base (includes CUDA and PyTorch)
FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime

# Set a working directory in the container
WORKDIR /workspace

# Install Python dependencies
COPY requirements.txt .
RUN pip install --upgrade pip
RUN pip install -r requirements.txt

# Copy the training script
COPY app.py .

# Run the training script
CMD ["python", "app.py"]