gary-boon Claude Opus 4.5 commited on
Commit
a009a49
·
1 Parent(s): 6435a75

Try PyTorch nightly for GB10/sm_121 GPU support

Browse files

The NVIDIA GB10 (Grace Blackwell) requires sm_121 compute capability
which isn't in stable PyTorch yet. Nightly may have it.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Files changed (1) hide show
  1. docker/Dockerfile.spark +4 -4
docker/Dockerfile.spark CHANGED
@@ -13,11 +13,11 @@ RUN apt-get update && apt-get install -y \
13
 
14
  WORKDIR /app
15
 
16
- # Install PyTorch with CUDA support first (before other requirements)
17
- # Using CUDA 12.4 wheels which are compatible with CUDA 13.0 driver
18
- RUN pip3 install --no-cache-dir \
19
  torch \
20
- --index-url https://download.pytorch.org/whl/cu124
21
 
22
  # Copy requirements and install (excluding packages with ARM64 issues)
23
  COPY requirements.txt .
 
13
 
14
  WORKDIR /app
15
 
16
+ # Install PyTorch nightly with CUDA support (for newer GPU architectures like GB10/sm_121)
17
+ # Using nightly builds which may have support for newer compute capabilities
18
+ RUN pip3 install --no-cache-dir --pre \
19
  torch \
20
+ --index-url https://download.pytorch.org/whl/nightly/cu124
21
 
22
  # Copy requirements and install (excluding packages with ARM64 issues)
23
  COPY requirements.txt .