Attempt to fix torchvision::nms error by installing PyTorch via pip with official CUDA index

Dockerfile

@@ -1,5 +1,4 @@
 # Start from a clean NVIDIA CUDA base image.
-# This provides the necessary CUDA runtime and development tools.
 # Using 12.4.0-devel-ubuntu22.04 to align with the CUDA version specified in your cosmos-predict1.yaml.
 FROM nvidia/cuda:12.4.0-devel-ubuntu22.04
 
@@ -19,7 +18,7 @@ WORKDIR /app
 RUN apt-get update && apt-get install -y --no-install-recommends \
     wget \
     git \
-    build-essential \ 
+    build-essential \
     libgl1-mesa-glx \
     # Clean up apt cache to reduce image size
     && rm -rf /var/lib/apt/lists/*
@@ -43,7 +42,7 @@ RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
 COPY . /app
 
 # Create the Conda environment named 'cosmos-predict1' using the provided YAML file.
-# This step will install all specified Python, PyTorch, CUDA, and pip dependencies.
+# This step will install all specified Python and pip dependencies (excluding PyTorch/TorchVision).
 RUN conda env create -f cosmos-predict1.yaml
 
 # Set the default Conda environment to be activated.
@@ -52,12 +51,24 @@ ENV CONDA_DEFAULT_ENV=cosmos-predict1
 # This ensures that executables (like python, pip, uvicorn) from this environment are found.
 ENV PATH=$CONDA_DIR/envs/cosmos-predict1/bin:$PATH
 
+# --- NEW: Install PyTorch and TorchVision via pip with specific CUDA index ---
+# This is crucial for ensuring the correct CUDA-enabled builds are installed.
+# We'll use CUDA 12.1 as it's commonly compatible with 12.4+ drivers.
+RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
+    conda activate cosmos-predict1 && \
+    pip install --no-cache-dir \
+        torch==2.3.1 \
+        torchvision==0.18.1 \
+        torchaudio==2.3.1 \
+        --index-url https://download.pytorch.org/whl/cu121
+# --- END NEW ---
+
 # --- Verification Steps ---
 RUN echo "Verifying Python and Conda installations..."
 RUN python --version
 RUN conda env list
 RUN echo "Verifying PyTorch and CUDA availability..."
-# REVISED AGAIN: Use a heredoc for multi-line Python code
+# Use a heredoc for multi-line Python code
 RUN conda run -n cosmos-predict1 python <<EOF
 import torch
 print('PyTorch Version: ' + torch.__version__)

cosmos-predict1.yaml

@@ -1,7 +1,6 @@
 name: cosmos-predict1
 channels:
-  - pytorch
-  - nvidia
+  # Keep these channels for other dependencies, but PyTorch will be installed via pip
   - conda-forge
   - defaults
 dependencies:
@@ -11,13 +10,13 @@ dependencies:
   - ninja
   - gcc=12.4.0
   - gxx=12.4.0
-
-  - pytorch=2.3.1
-  - torchvision=0.18.1
-  - torchaudio=2.3.1
-  - pytorch-cuda=12.4
-  - cudnn
-  - libcublas
+  # Remove all PyTorch/TorchVision/CUDA related entries from Conda dependencies
+  # - pytorch=2.3.1
+  # - torchvision=0.18.1
+  # - torchaudio=2.3.1
+  # - pytorch-cuda=12.1
+  # - cudnn
+  # - libcublas
 
   - pip:
       - -r file:///app/gui/requirements.txt
@@ -29,4 +28,5 @@ dependencies:
       - diffusers
       - megatron
       - IPython
-      - megatron.core
+      - megatron.core
+      # PyTorch and TorchVision will be installed separately in the Dockerfile