From 582ca44501f5c3e848937e765ec2281813447ee3 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Tue, 14 Apr 2026 21:33:35 +0800
Subject: [PATCH] Cache jitify result.

---
 containers/dockerfile/Dockerfile.gpu | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/containers/dockerfile/Dockerfile.gpu b/containers/dockerfile/Dockerfile.gpu
index 0fb06e6..104f308 100644
--- a/containers/dockerfile/Dockerfile.gpu
+++ b/containers/dockerfile/Dockerfile.gpu
@@ -54,6 +54,10 @@ RUN set -ex; \
     chmod +x /usr/local/bin/gosu && \
     gosu nobody true
 
+RUN \
+    source activate gpu_test && \
+    python -c 'from cupy.cuda import jitify; jitify._init_module()'
+
 # Default entry-point to use if running locally
 # It will preserve attributes of created files
 COPY entrypoint.sh /scripts/