diff --git a/GPU_puzzlers.ipynb b/GPU_puzzlers.ipynb
index 034ee2c..b4a98c4 100644
--- a/GPU_puzzlers.ipynb
+++ b/GPU_puzzlers.ipynb
@@ -127,6 +127,7 @@
         }
       ],
       "source": [
+        "!pip install -qqq numba-cuda==0.4.0\n",
         "!pip install -qqq git+https://github.com/chalk-diagrams/planar git+https://github.com/danoneata/chalk@srush-patch-1\n",
         "!wget -q https://github.com/srush/GPU-Puzzles/raw/main/robot.png https://github.com/srush/GPU-Puzzles/raw/main/lib.py"
       ]
@@ -6543,4 +6544,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 5
-}
\ No newline at end of file
+}
diff --git a/lib.py b/lib.py
index ce09e11..90e8da0 100644
--- a/lib.py
+++ b/lib.py
@@ -6,10 +6,13 @@
 from dataclasses import dataclass
 from typing import List, Any
 from collections import Counter
-from numba import cuda
+from numba import cuda, config
 import numba
 import random
 
+config.CUDA_ENABLE_PYNVJITLINK = 1
+config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
+
 @dataclass
 class ScalarHistory:
     last_fn: str