File size: 336 Bytes
d766458 | 1 2 3 4 5 6 7 8 9 10 | import os,jax
# disable triton_gemm for jax versions > 0.3
if int(jax.__version__.split(".")[1]) > 3:
os.environ["XLA_FLAGS"] = "--xla_gpu_enable_triton_gemm=false"
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
from colabdesign.shared.utils import clear_mem
from .model import mk_mpnn_model |