remat policy = None

Author: prishajain1

src/maxdiffusion/generate_ltx2.py

@@ -14,6 +14,7 @@
 
 from typing import Sequence
 import jax
+import jax.numpy as jnp
 import time
 import os
 import subprocess