.

Dockerfile

@@ -1,15 +1,29 @@
-# ベースイメージとしてPython 3.9を使用
-FROM python:3.9-slim
+# ベースイメージとしてCUDA対応のPythonを使用
+FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
 
 # 作業ディレクトリを設定
 WORKDIR /app
 
+# 必要な依存ライブラリをインストール
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3-pip \
+    python3-dev \
+    build-essential \
+    cuda-command-line-tools-12-1 \
+    cuda-nvcc-12-1 \
+    libcudnn8=8.9.0.*-1+cuda12.1 \
+    libcudnn8-dev=8.9.0.*-1+cuda12.1 \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
 # 必要なPythonライブラリをインストールするための依存ファイルをコピー
 COPY requirements.txt /app/requirements.txt
-
-# 必要なPythonパッケージをインストール
 RUN pip install --no-cache-dir -r requirements.txt
 
+# 環境変数を設定
+ENV PATH="/usr/local/cuda/bin:${PATH}"
+ENV XLA_FLAGS="--xla_gpu_cuda_data_dir=/usr/local/cuda"
+
 # アプリケーションコードをコンテナにコピー
 COPY . /app