bugfix for latest jax

sokrypton · Oct 9, 2022 · c90c4d3 · c90c4d3
1 parent 9494856
commit c90c4d3
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/colabdesign/shared/utils.py b/colabdesign/shared/utils.py
@@ -19,7 +19,10 @@ def clear_mem():
       for obj_name in dir(module):
         obj = getattr(module, obj_name)
         if hasattr(obj, "cache_clear"):
-          obj.cache_clear()
+          try:
+            obj.cache_clear()
+          except:
+            pass
   gc.collect()
 
 def update_dict(D, *args, **kwargs):
@@ -105,4 +108,4 @@ def softmax(x, axis=-1):
   return x / x.sum(axis,keepdims=True)
 
 def categorical(p):
-  return (p.cumsum(-1) >= np.random.uniform(size=p.shape[:-1])[..., None]).argmax(-1)
+  return (p.cumsum(-1) >= np.random.uniform(size=p.shape[:-1])[..., None]).argmax(-1)