Update Dockerfile

Dockerfile

@@ -24,9 +24,15 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     pip install scikit-learn==1.3.2 && \
     pip install matplotlib==3.7.2
 
-# --- Step 5b: Install PyTorch stack
+# --- Step 5b: Install PyTorch stack (one-by-one to save RAM)
 RUN --mount=type=cache,target=/root/.cache/pip \
-    pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
+    pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu118
+
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118
+
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
 
 # --- Step 5c: Install Hugging Face stack
 RUN --mount=type=cache,target=/root/.cache/pip \