File size: 427 Bytes
d766458 | 1 2 3 4 5 6 7 8 9 10 11 12 13 | 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 colabdesign.af.model import mk_af_model
# backward compatability
mk_design_model = mk_afdesign_model = mk_af_model |